From 94eb94e167fc2eb722c1526ca16e47254a038082 Mon Sep 17 00:00:00 2001
From: Savely Krendelhoff <alicecandyom@proton.me>
Date: Fri, 22 Aug 2025 21:32:46 +0700
Subject: [PATCH 1/7] [PHASE-5] Implement application layer

---
 .mockery.yaml                                 |   3 +
 internal/application/application.go           |  88 +++++
 internal/application/application_test.go      | 324 ++++++++++++++++++
 .../application/mocks/mock_wisdom_service.go  | 203 +++++++++++
 4 files changed, 618 insertions(+)
 create mode 100644 internal/application/application.go
 create mode 100644 internal/application/application_test.go
 create mode 100644 internal/application/mocks/mock_wisdom_service.go

diff --git a/.mockery.yaml b/.mockery.yaml
index 877ad49..6105c3d 100644
--- a/.mockery.yaml
+++ b/.mockery.yaml
@@ -9,3 +9,6 @@ packages:
       QuoteService:
       ChallengeGenerator:
       ChallengeVerifier:
+  hash-of-wisdom/internal/controller:
+    interfaces:
+      WisdomService:
diff --git a/internal/application/application.go b/internal/application/application.go
new file mode 100644
index 0000000..3ea5722
--- /dev/null
+++ b/internal/application/application.go
@@ -0,0 +1,88 @@
+package application
+
+import (
+	"context"
+	"fmt"
+	"io"
+
+	"hash-of-wisdom/internal/pow/challenge"
+	"hash-of-wisdom/internal/protocol"
+	"hash-of-wisdom/internal/quotes"
+)
+
+// Response represents an encodable response that can write itself to a connection
+type Response interface {
+	Encode(w io.Writer) error
+}
+
+// WisdomService defines the interface for the wisdom service
+type WisdomService interface {
+	GenerateChallenge(ctx context.Context, resource string) (*challenge.Challenge, error)
+	VerifySolution(ctx context.Context, solution *challenge.Solution) error
+	GetQuote(ctx context.Context) (*quotes.Quote, error)
+}
+
+// WisdomApplication handles the Word of Wisdom application logic
+type WisdomApplication struct {
+	wisdomService WisdomService
+}
+
+// NewWisdomApplication creates a new wisdom application handler
+func NewWisdomApplication(wisdomService WisdomService) *WisdomApplication {
+	return &WisdomApplication{
+		wisdomService: wisdomService,
+	}
+}
+
+// HandleMessage processes a protocol message and returns an encodable response
+func (a *WisdomApplication) HandleMessage(ctx context.Context, msg *protocol.Message) (Response, error) {
+	switch msg.Type {
+	case protocol.ChallengeRequestType:
+		return a.handleChallengeRequest(ctx)
+	case protocol.SolutionRequestType:
+		return a.handleSolutionRequest(ctx, msg)
+	default:
+		return &protocol.ErrorResponse{
+			Code:    protocol.ErrMalformedMessage,
+			Message: fmt.Sprintf("unsupported message type: 0x%02x", msg.Type),
+		}, nil
+	}
+}
+
+// handleChallengeRequest processes challenge requests
+func (a *WisdomApplication) handleChallengeRequest(ctx context.Context) (Response, error) {
+	challenge, err := a.wisdomService.GenerateChallenge(ctx, "quotes")
+	if err != nil {
+		return &protocol.ErrorResponse{Code: protocol.ErrServerError, Message: "Contact administrator"}, nil
+	}
+
+	return &protocol.ChallengeResponse{Challenge: challenge}, nil
+}
+
+// handleSolutionRequest processes solution requests
+func (a *WisdomApplication) handleSolutionRequest(ctx context.Context, msg *protocol.Message) (Response, error) {
+	// Parse solution request
+	var solutionReq protocol.SolutionRequest
+	if err := solutionReq.Decode(msg.PayloadStream); err != nil {
+		return &protocol.ErrorResponse{Code: protocol.ErrMalformedMessage, Message: "invalid solution format"}, nil
+	}
+
+	// Create solution object
+	solution := &challenge.Solution{
+		Challenge: solutionReq.Challenge,
+		Nonce:     solutionReq.Nonce,
+	}
+
+	// Verify solution
+	if err := a.wisdomService.VerifySolution(ctx, solution); err != nil {
+		return &protocol.ErrorResponse{Code: protocol.ErrInvalidSolution, Message: "solution verification failed"}, nil
+	}
+
+	// Get quote
+	quote, err := a.wisdomService.GetQuote(ctx)
+	if err != nil {
+		return &protocol.ErrorResponse{Code: protocol.ErrServerError, Message: "Contact administrator"}, nil
+	}
+
+	return &protocol.SolutionResponse{Quote: quote}, nil
+}
diff --git a/internal/application/application_test.go b/internal/application/application_test.go
new file mode 100644
index 0000000..6017a28
--- /dev/null
+++ b/internal/application/application_test.go
@@ -0,0 +1,324 @@
+package application
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"testing"
+
+	"hash-of-wisdom/internal/application/mocks"
+	"hash-of-wisdom/internal/pow/challenge"
+	"hash-of-wisdom/internal/protocol"
+	"hash-of-wisdom/internal/quotes"
+	"hash-of-wisdom/internal/service"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewWisdomApplication(t *testing.T) {
+	mockService := mocks.NewMockWisdomService(t)
+	app := NewWisdomApplication(mockService)
+
+	assert.NotNil(t, app)
+	assert.Equal(t, mockService, app.wisdomService)
+}
+
+func TestWisdomApplication_HandleMessage_UnsupportedType(t *testing.T) {
+	mockService := mocks.NewMockWisdomService(t)
+	app := NewWisdomApplication(mockService)
+
+	ctx := context.Background()
+	msg := &protocol.Message{
+		Type:           protocol.MessageType(0xFF), // Invalid type
+		PayloadLength:  0,
+		PayloadStream:  nil,
+	}
+
+	response, err := app.HandleMessage(ctx, msg)
+	require.NoError(t, err)
+
+	// Type assert to ErrorResponse
+	errorResponse, ok := response.(*protocol.ErrorResponse)
+	require.True(t, ok, "Expected ErrorResponse")
+	assert.Equal(t, protocol.ErrMalformedMessage, errorResponse.Code)
+	assert.Contains(t, errorResponse.Message, "unsupported message type: 0xff")
+}
+
+func TestWisdomApplication_HandleChallengeRequest_Success(t *testing.T) {
+	mockService := mocks.NewMockWisdomService(t)
+	app := NewWisdomApplication(mockService)
+
+	// Mock successful challenge generation
+	testChallenge := &challenge.Challenge{
+		Resource:   "quotes",
+		Timestamp:  12345,
+		Difficulty: 4,
+		Random:     []byte("test"),
+		HMAC:       []byte("signature"),
+	}
+
+	mockService.On("GenerateChallenge", mock.Anything, "quotes").Return(testChallenge, nil)
+
+	ctx := context.Background()
+	msg := &protocol.Message{
+		Type:           protocol.ChallengeRequestType,
+		PayloadLength:  0,
+		PayloadStream:  nil,
+	}
+
+	response, err := app.HandleMessage(ctx, msg)
+	require.NoError(t, err)
+
+	// Type assert to ChallengeResponse
+	challengeResponse, ok := response.(*protocol.ChallengeResponse)
+	require.True(t, ok, "Expected ChallengeResponse")
+	assert.Equal(t, testChallenge, challengeResponse.Challenge)
+
+	mockService.AssertExpectations(t)
+}
+
+func TestWisdomApplication_HandleChallengeRequest_ServiceError(t *testing.T) {
+	mockService := mocks.NewMockWisdomService(t)
+	app := NewWisdomApplication(mockService)
+
+	// Mock service error
+	mockService.On("GenerateChallenge", mock.Anything, "quotes").Return(nil, errors.New("service error"))
+
+	ctx := context.Background()
+	msg := &protocol.Message{
+		Type:           protocol.ChallengeRequestType,
+		PayloadLength:  0,
+		PayloadStream:  nil,
+	}
+
+	response, err := app.HandleMessage(ctx, msg)
+	require.NoError(t, err)
+
+	// Type assert to ErrorResponse
+	errorResponse, ok := response.(*protocol.ErrorResponse)
+	require.True(t, ok, "Expected ErrorResponse")
+	assert.Equal(t, protocol.ErrServerError, errorResponse.Code)
+	assert.Equal(t, "Contact administrator", errorResponse.Message)
+
+	mockService.AssertExpectations(t)
+}
+
+func TestWisdomApplication_HandleSolutionRequest_Success(t *testing.T) {
+	mockService := mocks.NewMockWisdomService(t)
+	app := NewWisdomApplication(mockService)
+
+	// Create test solution request
+	testChallenge := challenge.Challenge{
+		Resource:   "quotes",
+		Timestamp:  12345,
+		Difficulty: 4,
+		Random:     []byte("test"),
+		HMAC:       []byte("signature"),
+	}
+
+	solutionPayload := protocol.SolutionRequest{
+		Challenge: testChallenge,
+		Nonce:     12345,
+	}
+
+	payloadJSON, err := json.Marshal(solutionPayload)
+	require.NoError(t, err)
+
+	testQuote := &quotes.Quote{
+		Text:   "Test quote",
+		Author: "Test Author",
+	}
+
+	// Mock successful verification and quote retrieval
+	mockService.On("VerifySolution", mock.Anything, mock.AnythingOfType("*challenge.Solution")).Return(nil)
+	mockService.On("GetQuote", mock.Anything).Return(testQuote, nil)
+
+	ctx := context.Background()
+	msg := &protocol.Message{
+		Type:           protocol.SolutionRequestType,
+		PayloadLength:  uint32(len(payloadJSON)),
+		PayloadStream:  bytes.NewReader(payloadJSON),
+	}
+
+	response, err := app.HandleMessage(ctx, msg)
+	require.NoError(t, err)
+
+	// Type assert to SolutionResponse
+	solutionResponse, ok := response.(*protocol.SolutionResponse)
+	require.True(t, ok, "Expected SolutionResponse")
+	assert.Equal(t, testQuote, solutionResponse.Quote)
+
+	mockService.AssertExpectations(t)
+}
+
+func TestWisdomApplication_HandleSolutionRequest_InvalidJSON(t *testing.T) {
+	mockService := mocks.NewMockWisdomService(t)
+	app := NewWisdomApplication(mockService)
+
+	invalidJSON := []byte("invalid json")
+	ctx := context.Background()
+	msg := &protocol.Message{
+		Type:           protocol.SolutionRequestType,
+		PayloadLength:  uint32(len(invalidJSON)),
+		PayloadStream:  bytes.NewReader(invalidJSON),
+	}
+
+	response, err := app.HandleMessage(ctx, msg)
+	require.NoError(t, err)
+
+	// Type assert to ErrorResponse
+	errorResponse, ok := response.(*protocol.ErrorResponse)
+	require.True(t, ok, "Expected ErrorResponse")
+	assert.Equal(t, protocol.ErrMalformedMessage, errorResponse.Code)
+	assert.Equal(t, "invalid solution format", errorResponse.Message)
+
+	mockService.AssertNotCalled(t, "VerifySolution")
+	mockService.AssertNotCalled(t, "GetQuote")
+}
+
+func TestWisdomApplication_HandleSolutionRequest_VerificationFailed(t *testing.T) {
+	mockService := mocks.NewMockWisdomService(t)
+	app := NewWisdomApplication(mockService)
+
+	// Create test solution request
+	testChallenge := challenge.Challenge{
+		Resource:   "quotes",
+		Timestamp:  12345,
+		Difficulty: 4,
+		Random:     []byte("test"),
+		HMAC:       []byte("signature"),
+	}
+
+	solutionPayload := protocol.SolutionRequest{
+		Challenge: testChallenge,
+		Nonce:     12345,
+	}
+
+	payloadJSON, err := json.Marshal(solutionPayload)
+	require.NoError(t, err)
+
+	// Mock verification failure
+	mockService.On("VerifySolution", mock.Anything, mock.AnythingOfType("*challenge.Solution")).Return(service.ErrInvalidSolution)
+
+	ctx := context.Background()
+	msg := &protocol.Message{
+		Type:           protocol.SolutionRequestType,
+		PayloadLength:  uint32(len(payloadJSON)),
+		PayloadStream:  bytes.NewReader(payloadJSON),
+	}
+
+	response, err := app.HandleMessage(ctx, msg)
+	require.NoError(t, err)
+
+	// Type assert to ErrorResponse
+	errorResponse, ok := response.(*protocol.ErrorResponse)
+	require.True(t, ok, "Expected ErrorResponse")
+	assert.Equal(t, protocol.ErrInvalidSolution, errorResponse.Code)
+	assert.Equal(t, "solution verification failed", errorResponse.Message)
+
+	mockService.AssertExpectations(t)
+	mockService.AssertNotCalled(t, "GetQuote")
+}
+
+func TestWisdomApplication_HandleSolutionRequest_QuoteServiceError(t *testing.T) {
+	mockService := mocks.NewMockWisdomService(t)
+	app := NewWisdomApplication(mockService)
+
+	// Create test solution request
+	testChallenge := challenge.Challenge{
+		Resource:   "quotes",
+		Timestamp:  12345,
+		Difficulty: 4,
+		Random:     []byte("test"),
+		HMAC:       []byte("signature"),
+	}
+
+	solutionPayload := protocol.SolutionRequest{
+		Challenge: testChallenge,
+		Nonce:     12345,
+	}
+
+	payloadJSON, err := json.Marshal(solutionPayload)
+	require.NoError(t, err)
+
+	// Mock successful verification but quote service error
+	mockService.On("VerifySolution", mock.Anything, mock.AnythingOfType("*challenge.Solution")).Return(nil)
+	mockService.On("GetQuote", mock.Anything).Return(nil, errors.New("quote service error"))
+
+	ctx := context.Background()
+	msg := &protocol.Message{
+		Type:           protocol.SolutionRequestType,
+		PayloadLength:  uint32(len(payloadJSON)),
+		PayloadStream:  bytes.NewReader(payloadJSON),
+	}
+
+	response, err := app.HandleMessage(ctx, msg)
+	require.NoError(t, err)
+
+	// Type assert to ErrorResponse
+	errorResponse, ok := response.(*protocol.ErrorResponse)
+	require.True(t, ok, "Expected ErrorResponse")
+	assert.Equal(t, protocol.ErrServerError, errorResponse.Code)
+	assert.Equal(t, "Contact administrator", errorResponse.Message)
+
+	mockService.AssertExpectations(t)
+}
+
+func TestWisdomApplication_HandleMessage_ContextCancellation(t *testing.T) {
+	mockService := mocks.NewMockWisdomService(t)
+	app := NewWisdomApplication(mockService)
+
+	// Create cancelled context
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	// Mock service to respect context cancellation
+	mockService.On("GenerateChallenge", mock.Anything, "quotes").Return(nil, context.Canceled)
+
+	msg := &protocol.Message{
+		Type:           protocol.ChallengeRequestType,
+		PayloadLength:  0,
+		PayloadStream:  nil,
+	}
+
+	response, err := app.HandleMessage(ctx, msg)
+	require.NoError(t, err)
+
+	// Type assert to ErrorResponse
+	errorResponse, ok := response.(*protocol.ErrorResponse)
+	require.True(t, ok, "Expected ErrorResponse")
+	assert.Equal(t, protocol.ErrServerError, errorResponse.Code)
+
+	mockService.AssertExpectations(t)
+}
+
+func TestResponseEncoding(t *testing.T) {
+	// Test ChallengeResponse encoding produces valid binary format
+	testChallenge := &challenge.Challenge{
+		Resource:   "quotes",
+		Timestamp:  12345,
+		Difficulty: 4,
+		Random:     []byte("test"),
+		HMAC:       []byte("signature"),
+	}
+
+	challengeResponse := &protocol.ChallengeResponse{Challenge: testChallenge}
+	var buf bytes.Buffer
+	err := challengeResponse.Encode(&buf)
+	require.NoError(t, err)
+
+	// Verify binary format
+	data := buf.Bytes()
+	assert.GreaterOrEqual(t, len(data), 5) // At least header size
+
+	// Check message type
+	assert.Equal(t, byte(protocol.ChallengeResponseType), data[0])
+
+	// Check payload contains expected data
+	payload := string(data[5:]) // Skip header
+	assert.Contains(t, payload, "quotes")
+	assert.Contains(t, payload, "12345")
+}
diff --git a/internal/application/mocks/mock_wisdom_service.go b/internal/application/mocks/mock_wisdom_service.go
new file mode 100644
index 0000000..eb039bb
--- /dev/null
+++ b/internal/application/mocks/mock_wisdom_service.go
@@ -0,0 +1,203 @@
+// Code generated by mockery v2.53.5. DO NOT EDIT.
+
+package mocks
+
+import (
+	context "context"
+	challenge "hash-of-wisdom/internal/pow/challenge"
+
+	mock "github.com/stretchr/testify/mock"
+
+	quotes "hash-of-wisdom/internal/quotes"
+)
+
+// MockWisdomService is an autogenerated mock type for the WisdomService type
+type MockWisdomService struct {
+	mock.Mock
+}
+
+type MockWisdomService_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockWisdomService) EXPECT() *MockWisdomService_Expecter {
+	return &MockWisdomService_Expecter{mock: &_m.Mock}
+}
+
+// GenerateChallenge provides a mock function with given fields: ctx, resource
+func (_m *MockWisdomService) GenerateChallenge(ctx context.Context, resource string) (*challenge.Challenge, error) {
+	ret := _m.Called(ctx, resource)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GenerateChallenge")
+	}
+
+	var r0 *challenge.Challenge
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, string) (*challenge.Challenge, error)); ok {
+		return rf(ctx, resource)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, string) *challenge.Challenge); ok {
+		r0 = rf(ctx, resource)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*challenge.Challenge)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = rf(ctx, resource)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockWisdomService_GenerateChallenge_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GenerateChallenge'
+type MockWisdomService_GenerateChallenge_Call struct {
+	*mock.Call
+}
+
+// GenerateChallenge is a helper method to define mock.On call
+//   - ctx context.Context
+//   - resource string
+func (_e *MockWisdomService_Expecter) GenerateChallenge(ctx interface{}, resource interface{}) *MockWisdomService_GenerateChallenge_Call {
+	return &MockWisdomService_GenerateChallenge_Call{Call: _e.mock.On("GenerateChallenge", ctx, resource)}
+}
+
+func (_c *MockWisdomService_GenerateChallenge_Call) Run(run func(ctx context.Context, resource string)) *MockWisdomService_GenerateChallenge_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(string))
+	})
+	return _c
+}
+
+func (_c *MockWisdomService_GenerateChallenge_Call) Return(_a0 *challenge.Challenge, _a1 error) *MockWisdomService_GenerateChallenge_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockWisdomService_GenerateChallenge_Call) RunAndReturn(run func(context.Context, string) (*challenge.Challenge, error)) *MockWisdomService_GenerateChallenge_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// GetQuote provides a mock function with given fields: ctx
+func (_m *MockWisdomService) GetQuote(ctx context.Context) (*quotes.Quote, error) {
+	ret := _m.Called(ctx)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetQuote")
+	}
+
+	var r0 *quotes.Quote
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context) (*quotes.Quote, error)); ok {
+		return rf(ctx)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context) *quotes.Quote); ok {
+		r0 = rf(ctx)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*quotes.Quote)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = rf(ctx)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockWisdomService_GetQuote_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetQuote'
+type MockWisdomService_GetQuote_Call struct {
+	*mock.Call
+}
+
+// GetQuote is a helper method to define mock.On call
+//   - ctx context.Context
+func (_e *MockWisdomService_Expecter) GetQuote(ctx interface{}) *MockWisdomService_GetQuote_Call {
+	return &MockWisdomService_GetQuote_Call{Call: _e.mock.On("GetQuote", ctx)}
+}
+
+func (_c *MockWisdomService_GetQuote_Call) Run(run func(ctx context.Context)) *MockWisdomService_GetQuote_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context))
+	})
+	return _c
+}
+
+func (_c *MockWisdomService_GetQuote_Call) Return(_a0 *quotes.Quote, _a1 error) *MockWisdomService_GetQuote_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockWisdomService_GetQuote_Call) RunAndReturn(run func(context.Context) (*quotes.Quote, error)) *MockWisdomService_GetQuote_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// VerifySolution provides a mock function with given fields: ctx, solution
+func (_m *MockWisdomService) VerifySolution(ctx context.Context, solution *challenge.Solution) error {
+	ret := _m.Called(ctx, solution)
+
+	if len(ret) == 0 {
+		panic("no return value specified for VerifySolution")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, *challenge.Solution) error); ok {
+		r0 = rf(ctx, solution)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockWisdomService_VerifySolution_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'VerifySolution'
+type MockWisdomService_VerifySolution_Call struct {
+	*mock.Call
+}
+
+// VerifySolution is a helper method to define mock.On call
+//   - ctx context.Context
+//   - solution *challenge.Solution
+func (_e *MockWisdomService_Expecter) VerifySolution(ctx interface{}, solution interface{}) *MockWisdomService_VerifySolution_Call {
+	return &MockWisdomService_VerifySolution_Call{Call: _e.mock.On("VerifySolution", ctx, solution)}
+}
+
+func (_c *MockWisdomService_VerifySolution_Call) Run(run func(ctx context.Context, solution *challenge.Solution)) *MockWisdomService_VerifySolution_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*challenge.Solution))
+	})
+	return _c
+}
+
+func (_c *MockWisdomService_VerifySolution_Call) Return(_a0 error) *MockWisdomService_VerifySolution_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockWisdomService_VerifySolution_Call) RunAndReturn(run func(context.Context, *challenge.Solution) error) *MockWisdomService_VerifySolution_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockWisdomService creates a new instance of MockWisdomService. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockWisdomService(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockWisdomService {
+	mock := &MockWisdomService{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
-- 
2.44.1


From d12de089a0bc8c837f7a3fde95aba58b4d8b3343 Mon Sep 17 00:00:00 2001
From: Savely Krendelhoff <alicecandyom@proton.me>
Date: Sat, 23 Aug 2025 12:04:38 +0700
Subject: [PATCH 2/7] [PHASE-5] Delegate encoding to the objects themselves

---
 internal/protocol/codec.go      |  43 +-----
 internal/protocol/codec_test.go | 265 +++++++++++---------------------
 internal/protocol/responses.go  |  82 ++++++++++
 3 files changed, 173 insertions(+), 217 deletions(-)
 create mode 100644 internal/protocol/responses.go

diff --git a/internal/protocol/codec.go b/internal/protocol/codec.go
index e5d83cd..d4c0d20 100644
--- a/internal/protocol/codec.go
+++ b/internal/protocol/codec.go
@@ -15,38 +15,6 @@ func NewCodec() *Codec {
 	return &Codec{}
 }
 
-// Encode writes a message to the writer using the protocol format
-func (c *Codec) Encode(w io.Writer, msg *Message) error {
-	if msg == nil {
-		return fmt.Errorf("message cannot be nil")
-	}
-
-	// Validate payload size
-	if len(msg.Payload) > MaxPayloadSize {
-		return fmt.Errorf("payload size %d exceeds maximum %d", len(msg.Payload), MaxPayloadSize)
-	}
-
-	// Write message type (1 byte)
-	if err := binary.Write(w, binary.BigEndian, msg.Type); err != nil {
-		return fmt.Errorf("failed to write message type: %w", err)
-	}
-
-	// Write payload length (4 bytes, big-endian)
-	payloadLength := uint32(len(msg.Payload))
-	if err := binary.Write(w, binary.BigEndian, payloadLength); err != nil {
-		return fmt.Errorf("failed to write payload length: %w", err)
-	}
-
-	// Write payload if present
-	if len(msg.Payload) > 0 {
-		if _, err := w.Write(msg.Payload); err != nil {
-			return fmt.Errorf("failed to write payload: %w", err)
-		}
-	}
-
-	return nil
-}
-
 // Decode reads a message from the reader using the protocol format
 func (c *Codec) Decode(r io.Reader) (*Message, error) {
 	// Read message type (1 byte)
@@ -78,12 +46,9 @@ func (c *Codec) Decode(r io.Reader) (*Message, error) {
 	var payload []byte
 	if payloadLength > 0 {
 		payload = make([]byte, payloadLength)
-		// Use LimitReader to ensure we don't read more than payloadLength bytes
-		// even if the underlying reader has more data available
-		limitedReader := io.LimitReader(r, int64(payloadLength))
-		// Note: ReadFull may block waiting for data. The connection handler
-		// MUST set appropriate read deadlines to prevent slowloris attacks
-		if _, err := io.ReadFull(limitedReader, payload); err != nil {
+		// ReadFull reads exactly payloadLength bytes
+		// The server MUST use LimitReader and set read deadlines to prevent attacks
+		if _, err := io.ReadFull(r, payload); err != nil {
 			return nil, fmt.Errorf("failed to read payload: %w", err)
 		}
 
@@ -103,7 +68,7 @@ func (c *Codec) Decode(r io.Reader) (*Message, error) {
 // isValidMessageType checks if the message type is defined in the protocol
 func isValidMessageType(msgType MessageType) bool {
 	switch msgType {
-	case ChallengeRequest, ChallengeResponse, SolutionRequest, QuoteResponse, ErrorResponse:
+	case ChallengeRequestType, ChallengeResponseType, SolutionRequestType, QuoteResponseType, ErrorResponseType:
 		return true
 	default:
 		return false
diff --git a/internal/protocol/codec_test.go b/internal/protocol/codec_test.go
index 6e5d6d5..f0eb749 100644
--- a/internal/protocol/codec_test.go
+++ b/internal/protocol/codec_test.go
@@ -2,102 +2,18 @@ package protocol
 
 import (
 	"bytes"
-	"encoding/json"
 	"io"
 	"testing"
 	"time"
 
+	"hash-of-wisdom/internal/pow/challenge"
+	"hash-of-wisdom/internal/quotes"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-func TestCodec_Encode_Decode(t *testing.T) {
-	codec := NewCodec()
-
-	tests := []struct {
-		name    string
-		message *Message
-	}{
-		{
-			name: "challenge request (empty payload)",
-			message: &Message{
-				Type:    ChallengeRequest,
-				Payload: []byte{},
-			},
-		},
-		{
-			name: "challenge response with payload",
-			message: &Message{
-				Type:    ChallengeResponse,
-				Payload: []byte(`{"challenge":{"timestamp":1640995200,"difficulty":4}}`),
-			},
-		},
-		{
-			name: "error response",
-			message: &Message{
-				Type:    ErrorResponse,
-				Payload: []byte(`{"code":"INVALID_SOLUTION","message":"Invalid nonce"}`),
-			},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			var buf bytes.Buffer
-
-			// Encode message
-			err := codec.Encode(&buf, tt.message)
-			require.NoError(t, err)
-
-			// Decode message
-			decoded, err := codec.Decode(&buf)
-			require.NoError(t, err)
-
-			assert.Equal(t, tt.message.Type, decoded.Type)
-			if len(tt.message.Payload) == 0 && len(decoded.Payload) == 0 {
-				// Both are empty (nil or empty slice)
-				assert.True(t, true)
-			} else {
-				assert.Equal(t, tt.message.Payload, decoded.Payload)
-			}
-		})
-	}
-}
-
-func TestCodec_Encode_Errors(t *testing.T) {
-	codec := NewCodec()
-
-	tests := []struct {
-		name    string
-		message *Message
-		wantErr string
-	}{
-		{
-			name:    "nil message",
-			message: nil,
-			wantErr: "message cannot be nil",
-		},
-		{
-			name: "payload too large",
-			message: &Message{
-				Type:    ChallengeRequest,
-				Payload: make([]byte, MaxPayloadSize+1),
-			},
-			wantErr: "payload size",
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			var buf bytes.Buffer
-			err := codec.Encode(&buf, tt.message)
-			assert.Error(t, err)
-			assert.Contains(t, err.Error(), tt.wantErr)
-		})
-	}
-}
-
-func TestCodec_Decode_Errors(t *testing.T) {
+func TestCodec_Decode(t *testing.T) {
 	codec := NewCodec()
 
 	tests := []struct {
@@ -147,96 +63,6 @@ func TestCodec_Decode_Errors(t *testing.T) {
 	}
 }
 
-
-func TestCodec_RoundTrip_RealPayloads(t *testing.T) {
-	codec := NewCodec()
-
-	t.Run("challenge response round trip", func(t *testing.T) {
-		original := &ChallengeResponsePayload{
-			Timestamp:  time.Now().Unix(),
-			Difficulty: 4,
-			Resource:   "quotes",
-			Random:     []byte("random123"),
-			HMAC:       []byte("hmac_signature"),
-		}
-
-		// Marshal payload
-		jsonData, err := json.Marshal(original)
-		require.NoError(t, err)
-
-		msg := &Message{
-			Type:    ChallengeResponse,
-			Payload: jsonData,
-		}
-
-		// Simulate network transmission
-		var buf bytes.Buffer
-		err = codec.Encode(&buf, msg)
-		require.NoError(t, err)
-
-		// Decode message
-		decoded, err := codec.Decode(&buf)
-		require.NoError(t, err)
-
-		// Unmarshal payload
-		var result ChallengeResponsePayload
-		err = json.Unmarshal(decoded.Payload, &result)
-		require.NoError(t, err)
-
-		assert.Equal(t, original.Timestamp, result.Timestamp)
-		assert.Equal(t, original.Difficulty, result.Difficulty)
-		assert.Equal(t, original.Resource, result.Resource)
-		assert.Equal(t, original.Random, result.Random)
-		assert.Equal(t, original.HMAC, result.HMAC)
-	})
-
-	t.Run("quote response round trip", func(t *testing.T) {
-		original := &QuoteResponsePayload{
-			Text:   "Test quote",
-			Author: "Test author",
-		}
-
-		// Marshal payload
-		jsonData, err := json.Marshal(original)
-		require.NoError(t, err)
-
-		msg := &Message{
-			Type:    QuoteResponse,
-			Payload: jsonData,
-		}
-
-		var buf bytes.Buffer
-		err = codec.Encode(&buf, msg)
-		require.NoError(t, err)
-
-		decoded, err := codec.Decode(&buf)
-		require.NoError(t, err)
-
-		var result QuoteResponsePayload
-		err = json.Unmarshal(decoded.Payload, &result)
-		require.NoError(t, err)
-
-		assert.Equal(t, original.Text, result.Text)
-		assert.Equal(t, original.Author, result.Author)
-	})
-}
-
-func TestCodec_WriteError_Handling(t *testing.T) {
-	codec := NewCodec()
-
-	// Create a writer that fails after a certain number of bytes
-	failAfter := 3
-	writer := &failingWriter{failAfter: failAfter}
-
-	msg := &Message{
-		Type:    ChallengeResponse,
-		Payload: []byte("test payload"),
-	}
-
-	err := codec.Encode(writer, msg)
-	assert.Error(t, err)
-}
-
 func TestCodec_ReadError_Handling(t *testing.T) {
 	codec := NewCodec()
 
@@ -251,6 +77,89 @@ func TestCodec_ReadError_Handling(t *testing.T) {
 	assert.Contains(t, err.Error(), "failed to read payload")
 }
 
+func TestChallengeResponse_Encode(t *testing.T) {
+	challenge := &challenge.Challenge{
+		Timestamp:  time.Now().Unix(),
+		Difficulty: 4,
+		Resource:   "quotes",
+		Random:     []byte("random123"),
+		HMAC:       []byte("hmac_signature"),
+	}
+
+	response := &ChallengeResponse{Challenge: challenge}
+
+	var buf bytes.Buffer
+	err := response.Encode(&buf)
+	require.NoError(t, err)
+
+	// Verify the encoded message can be decoded
+	codec := NewCodec()
+	decoded, err := codec.Decode(&buf)
+	require.NoError(t, err)
+
+	assert.Equal(t, ChallengeResponseType, decoded.Type)
+	assert.Contains(t, string(decoded.Payload), "quotes")
+	assert.Contains(t, string(decoded.Payload), "cmFuZG9tMTIz") // "random123" base64 encoded
+}
+
+func TestSolutionResponse_Encode(t *testing.T) {
+	quote := &quotes.Quote{
+		Text:   "Test quote",
+		Author: "Test author",
+	}
+
+	response := &SolutionResponse{Quote: quote}
+
+	var buf bytes.Buffer
+	err := response.Encode(&buf)
+	require.NoError(t, err)
+
+	// Verify the encoded message can be decoded
+	codec := NewCodec()
+	decoded, err := codec.Decode(&buf)
+	require.NoError(t, err)
+
+	assert.Equal(t, QuoteResponseType, decoded.Type)
+	assert.Contains(t, string(decoded.Payload), "Test quote")
+	assert.Contains(t, string(decoded.Payload), "Test author")
+}
+
+func TestErrorResponse_Encode(t *testing.T) {
+	errorResp := &ErrorResponse{
+		Code:       "INVALID_SOLUTION",
+		Message:    "The provided PoW solution is incorrect",
+		RetryAfter: 30,
+		Details:    map[string]string{"attempt": "1"},
+	}
+
+	var buf bytes.Buffer
+	err := errorResp.Encode(&buf)
+	require.NoError(t, err)
+
+	// Verify the encoded message can be decoded
+	codec := NewCodec()
+	decoded, err := codec.Decode(&buf)
+	require.NoError(t, err)
+
+	assert.Equal(t, ErrorResponseType, decoded.Type)
+	assert.Contains(t, string(decoded.Payload), "INVALID_SOLUTION")
+	assert.Contains(t, string(decoded.Payload), "The provided PoW solution is incorrect")
+	assert.Contains(t, string(decoded.Payload), "30")
+}
+
+func TestResponse_WriteError_Handling(t *testing.T) {
+	response := &ErrorResponse{
+		Code:    "TEST_ERROR",
+		Message: "Test message",
+	}
+
+	// Create a writer that fails immediately
+	writer := &failingWriter{failAfter: 1}
+
+	err := response.Encode(writer)
+	assert.Error(t, err)
+}
+
 // Helper functions and types for testing
 
 func encodeBigEndianUint32(val uint32) []byte {
diff --git a/internal/protocol/responses.go b/internal/protocol/responses.go
new file mode 100644
index 0000000..8521d4f
--- /dev/null
+++ b/internal/protocol/responses.go
@@ -0,0 +1,82 @@
+package protocol
+
+import (
+	"encoding/binary"
+	"encoding/json"
+	"fmt"
+	"io"
+
+	"hash-of-wisdom/internal/pow/challenge"
+	"hash-of-wisdom/internal/quotes"
+)
+
+// writeHeader writes the message type and payload length to the writer
+func writeHeader(w io.Writer, msgType MessageType, payloadLength uint32) error {
+	// Write message type (1 byte)
+	if err := binary.Write(w, binary.BigEndian, msgType); err != nil {
+		return fmt.Errorf("failed to write message type: %w", err)
+	}
+
+	// Write payload length (4 bytes, big-endian)
+	if err := binary.Write(w, binary.BigEndian, payloadLength); err != nil {
+		return fmt.Errorf("failed to write payload length: %w", err)
+	}
+
+	return nil
+}
+
+// encodeResponse is a helper function that encodes any response with the given message type
+func encodeResponse(w io.Writer, msgType MessageType, payload interface{}) error {
+	// Marshal to get exact payload size
+	payloadBytes, err := json.Marshal(payload)
+	if err != nil {
+		return fmt.Errorf("failed to encode payload: %w", err)
+	}
+
+	// Write header
+	if err := writeHeader(w, msgType, uint32(len(payloadBytes))); err != nil {
+		return err
+	}
+
+	// Write JSON payload directly to stream
+	if len(payloadBytes) > 0 {
+		if _, err := w.Write(payloadBytes); err != nil {
+			return fmt.Errorf("failed to write payload: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// ChallengeResponse represents a challenge response
+type ChallengeResponse struct {
+	Challenge *challenge.Challenge
+}
+
+// SolutionResponse represents a successful solution response (contains quote)
+type SolutionResponse struct {
+	Quote *quotes.Quote
+}
+
+// ErrorResponse represents an error response
+type ErrorResponse struct {
+	Code       string            `json:"code"`
+	Message    string            `json:"message"`
+	RetryAfter int               `json:"retry_after,omitempty"`
+	Details    map[string]string `json:"details,omitempty"`
+}
+
+// Encode writes the challenge response to the writer
+func (r *ChallengeResponse) Encode(w io.Writer) error {
+	return encodeResponse(w, ChallengeResponseType, r.Challenge)
+}
+
+// Encode writes the solution response to the writer
+func (r *SolutionResponse) Encode(w io.Writer) error {
+	return encodeResponse(w, QuoteResponseType, r.Quote)
+}
+
+// Encode writes the error response to the writer
+func (r *ErrorResponse) Encode(w io.Writer) error {
+	return encodeResponse(w, ErrorResponseType, r)
+}
-- 
2.44.1


From c7b930c9357bb89a62751738072c2d82e433dc47 Mon Sep 17 00:00:00 2001
From: Savely Krendelhoff <alicecandyom@proton.me>
Date: Sat, 23 Aug 2025 12:08:09 +0700
Subject: [PATCH 3/7] [PHASE-5] Move constants into separate file

---
 internal/protocol/constants.go | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)
 create mode 100644 internal/protocol/constants.go

diff --git a/internal/protocol/constants.go b/internal/protocol/constants.go
new file mode 100644
index 0000000..015853c
--- /dev/null
+++ b/internal/protocol/constants.go
@@ -0,0 +1,31 @@
+package protocol
+
+// MessageType represents the type of protocol message
+type MessageType byte
+
+const (
+	ChallengeRequestType  MessageType = 0x01
+	SolutionRequestType   MessageType = 0x03
+	// Response types (for responses.go)
+	ChallengeResponseType MessageType = 0x02
+	QuoteResponseType     MessageType = 0x04
+	ErrorResponseType     MessageType = 0x05
+)
+
+// Error codes as defined in protocol specification
+const (
+	ErrMalformedMessage    = "MALFORMED_MESSAGE"
+	ErrInvalidChallenge    = "INVALID_CHALLENGE"
+	ErrInvalidSolution     = "INVALID_SOLUTION"
+	ErrExpiredChallenge    = "EXPIRED_CHALLENGE"
+	ErrRateLimited         = "RATE_LIMITED"
+	ErrServerError         = "SERVER_ERROR"
+	ErrTooManyConnections  = "TOO_MANY_CONNECTIONS"
+	ErrDifficultyTooHigh   = "DIFFICULTY_TOO_HIGH"
+)
+
+// Protocol constants
+const (
+	MaxPayloadSize = 8 * 1024 // 8KB maximum payload size
+	HeaderSize     = 5        // 1 byte type + 4 bytes length
+)
-- 
2.44.1


From c147bc7fe4da0d6af3f40dd394c277a3c51d12aa Mon Sep 17 00:00:00 2001
From: Savely Krendelhoff <alicecandyom@proton.me>
Date: Sat, 23 Aug 2025 12:16:43 +0700
Subject: [PATCH 4/7] [PHASE-5] Move requests to separate file and implement
 decoding

---
 internal/protocol/requests.go | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)
 create mode 100644 internal/protocol/requests.go

diff --git a/internal/protocol/requests.go b/internal/protocol/requests.go
new file mode 100644
index 0000000..3961223
--- /dev/null
+++ b/internal/protocol/requests.go
@@ -0,0 +1,35 @@
+package protocol
+
+import (
+	"encoding/json"
+	"io"
+
+	"hash-of-wisdom/internal/pow/challenge"
+)
+
+// ChallengeRequest is empty (no payload for challenge requests)
+type ChallengeRequest struct{}
+
+// Decode reads a challenge request from the payload stream
+func (r *ChallengeRequest) Decode(stream io.Reader) error {
+	// Challenge requests have no payload to decode
+	return nil
+}
+
+// SolutionRequest contains the client's solution attempt
+type SolutionRequest struct {
+	Challenge challenge.Challenge `json:"challenge"`
+	Nonce     uint64              `json:"nonce"`
+}
+
+// Decode reads a solution request from the payload stream
+func (r *SolutionRequest) Decode(stream io.Reader) error {
+	if stream == nil {
+		// json.NewDecoder panics on nil reader
+		return io.EOF
+	}
+
+	// Parse JSON directly from stream
+	decoder := json.NewDecoder(stream)
+	return decoder.Decode(r)
+}
-- 
2.44.1


From 140099d6c2606ebf1cce773cc3d452b5d9cc24ab Mon Sep 17 00:00:00 2001
From: Savely Krendelhoff <alicecandyom@proton.me>
Date: Sat, 23 Aug 2025 12:18:45 +0700
Subject: [PATCH 5/7] [PHASE-5] Rework codec

---
 internal/protocol/codec.go                |  76 --------
 internal/protocol/codec_test.go           | 217 ----------------------
 internal/protocol/message_decoder.go      |  65 +++++++
 internal/protocol/message_decoder_test.go | 206 ++++++++++++++++++++
 internal/protocol/types.go                |  64 +------
 5 files changed, 276 insertions(+), 352 deletions(-)
 delete mode 100644 internal/protocol/codec.go
 delete mode 100644 internal/protocol/codec_test.go
 create mode 100644 internal/protocol/message_decoder.go
 create mode 100644 internal/protocol/message_decoder_test.go

diff --git a/internal/protocol/codec.go b/internal/protocol/codec.go
deleted file mode 100644
index d4c0d20..0000000
--- a/internal/protocol/codec.go
+++ /dev/null
@@ -1,76 +0,0 @@
-package protocol
-
-import (
-	"encoding/binary"
-	"fmt"
-	"io"
-	"unicode/utf8"
-)
-
-// Codec handles encoding and decoding of protocol messages
-type Codec struct{}
-
-// NewCodec creates a new protocol codec
-func NewCodec() *Codec {
-	return &Codec{}
-}
-
-// Decode reads a message from the reader using the protocol format
-func (c *Codec) Decode(r io.Reader) (*Message, error) {
-	// Read message type (1 byte)
-	var msgType MessageType
-	if err := binary.Read(r, binary.BigEndian, &msgType); err != nil {
-		if err == io.EOF {
-			return nil, err
-		}
-		return nil, fmt.Errorf("failed to read message type: %w", err)
-	}
-
-	// Validate message type
-	if !isValidMessageType(msgType) {
-		return nil, fmt.Errorf("invalid message type: 0x%02x", msgType)
-	}
-
-	// Read payload length (4 bytes, big-endian)
-	var payloadLength uint32
-	if err := binary.Read(r, binary.BigEndian, &payloadLength); err != nil {
-		return nil, fmt.Errorf("failed to read payload length: %w", err)
-	}
-
-	// Validate payload length
-	if payloadLength > MaxPayloadSize {
-		return nil, fmt.Errorf("payload length %d exceeds maximum %d", payloadLength, MaxPayloadSize)
-	}
-
-	// Read payload if present
-	var payload []byte
-	if payloadLength > 0 {
-		payload = make([]byte, payloadLength)
-		// ReadFull reads exactly payloadLength bytes
-		// The server MUST use LimitReader and set read deadlines to prevent attacks
-		if _, err := io.ReadFull(r, payload); err != nil {
-			return nil, fmt.Errorf("failed to read payload: %w", err)
-		}
-
-		// Validate payload is valid UTF-8
-		if !utf8.Valid(payload) {
-			return nil, fmt.Errorf("payload contains invalid UTF-8")
-		}
-	}
-
-	return &Message{
-		Type:    msgType,
-		Payload: payload,
-	}, nil
-}
-
-
-// isValidMessageType checks if the message type is defined in the protocol
-func isValidMessageType(msgType MessageType) bool {
-	switch msgType {
-	case ChallengeRequestType, ChallengeResponseType, SolutionRequestType, QuoteResponseType, ErrorResponseType:
-		return true
-	default:
-		return false
-	}
-}
diff --git a/internal/protocol/codec_test.go b/internal/protocol/codec_test.go
deleted file mode 100644
index f0eb749..0000000
--- a/internal/protocol/codec_test.go
+++ /dev/null
@@ -1,217 +0,0 @@
-package protocol
-
-import (
-	"bytes"
-	"io"
-	"testing"
-	"time"
-
-	"hash-of-wisdom/internal/pow/challenge"
-	"hash-of-wisdom/internal/quotes"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-)
-
-func TestCodec_Decode(t *testing.T) {
-	codec := NewCodec()
-
-	tests := []struct {
-		name    string
-		data    []byte
-		wantErr string
-	}{
-		{
-			name:    "empty data",
-			data:    []byte{},
-			wantErr: "EOF",
-		},
-		{
-			name:    "invalid message type",
-			data:    []byte{0xFF, 0x00, 0x00, 0x00, 0x00},
-			wantErr: "invalid message type",
-		},
-		{
-			name:    "incomplete header",
-			data:    []byte{0x01, 0x00, 0x00},
-			wantErr: "failed to read payload length",
-		},
-		{
-			name:    "payload too large",
-			data:    append([]byte{0x01}, encodeBigEndianUint32(MaxPayloadSize+1)...),
-			wantErr: "payload length",
-		},
-		{
-			name:    "incomplete payload",
-			data:    []byte{0x01, 0x00, 0x00, 0x00, 0x05, 0x01, 0x02},
-			wantErr: "failed to read payload",
-		},
-		{
-			name:    "invalid UTF-8 in payload",
-			data:    []byte{0x01, 0x00, 0x00, 0x00, 0x03, 0xFF, 0xFE, 0xFD},
-			wantErr: "invalid UTF-8",
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			buf := bytes.NewBuffer(tt.data)
-			_, err := codec.Decode(buf)
-			assert.Error(t, err)
-			assert.Contains(t, err.Error(), tt.wantErr)
-		})
-	}
-}
-
-func TestCodec_ReadError_Handling(t *testing.T) {
-	codec := NewCodec()
-
-	// Create a reader that fails after reading header
-	reader := &failingReader{
-		data:      []byte{0x01, 0x00, 0x00, 0x00, 0x05},
-		failAfter: 5,
-	}
-
-	_, err := codec.Decode(reader)
-	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "failed to read payload")
-}
-
-func TestChallengeResponse_Encode(t *testing.T) {
-	challenge := &challenge.Challenge{
-		Timestamp:  time.Now().Unix(),
-		Difficulty: 4,
-		Resource:   "quotes",
-		Random:     []byte("random123"),
-		HMAC:       []byte("hmac_signature"),
-	}
-
-	response := &ChallengeResponse{Challenge: challenge}
-
-	var buf bytes.Buffer
-	err := response.Encode(&buf)
-	require.NoError(t, err)
-
-	// Verify the encoded message can be decoded
-	codec := NewCodec()
-	decoded, err := codec.Decode(&buf)
-	require.NoError(t, err)
-
-	assert.Equal(t, ChallengeResponseType, decoded.Type)
-	assert.Contains(t, string(decoded.Payload), "quotes")
-	assert.Contains(t, string(decoded.Payload), "cmFuZG9tMTIz") // "random123" base64 encoded
-}
-
-func TestSolutionResponse_Encode(t *testing.T) {
-	quote := &quotes.Quote{
-		Text:   "Test quote",
-		Author: "Test author",
-	}
-
-	response := &SolutionResponse{Quote: quote}
-
-	var buf bytes.Buffer
-	err := response.Encode(&buf)
-	require.NoError(t, err)
-
-	// Verify the encoded message can be decoded
-	codec := NewCodec()
-	decoded, err := codec.Decode(&buf)
-	require.NoError(t, err)
-
-	assert.Equal(t, QuoteResponseType, decoded.Type)
-	assert.Contains(t, string(decoded.Payload), "Test quote")
-	assert.Contains(t, string(decoded.Payload), "Test author")
-}
-
-func TestErrorResponse_Encode(t *testing.T) {
-	errorResp := &ErrorResponse{
-		Code:       "INVALID_SOLUTION",
-		Message:    "The provided PoW solution is incorrect",
-		RetryAfter: 30,
-		Details:    map[string]string{"attempt": "1"},
-	}
-
-	var buf bytes.Buffer
-	err := errorResp.Encode(&buf)
-	require.NoError(t, err)
-
-	// Verify the encoded message can be decoded
-	codec := NewCodec()
-	decoded, err := codec.Decode(&buf)
-	require.NoError(t, err)
-
-	assert.Equal(t, ErrorResponseType, decoded.Type)
-	assert.Contains(t, string(decoded.Payload), "INVALID_SOLUTION")
-	assert.Contains(t, string(decoded.Payload), "The provided PoW solution is incorrect")
-	assert.Contains(t, string(decoded.Payload), "30")
-}
-
-func TestResponse_WriteError_Handling(t *testing.T) {
-	response := &ErrorResponse{
-		Code:    "TEST_ERROR",
-		Message: "Test message",
-	}
-
-	// Create a writer that fails immediately
-	writer := &failingWriter{failAfter: 1}
-
-	err := response.Encode(writer)
-	assert.Error(t, err)
-}
-
-// Helper functions and types for testing
-
-func encodeBigEndianUint32(val uint32) []byte {
-	return []byte{
-		byte(val >> 24),
-		byte(val >> 16),
-		byte(val >> 8),
-		byte(val),
-	}
-}
-
-type failingWriter struct {
-	written   int
-	failAfter int
-}
-
-func (w *failingWriter) Write(data []byte) (int, error) {
-	if w.written >= w.failAfter {
-		return 0, io.ErrShortWrite
-	}
-
-	remaining := w.failAfter - w.written
-	if len(data) <= remaining {
-		w.written += len(data)
-		return len(data), nil
-	}
-
-	w.written = w.failAfter
-	return remaining, io.ErrShortWrite
-}
-
-type failingReader struct {
-	data      []byte
-	pos       int
-	failAfter int
-}
-
-func (r *failingReader) Read(buf []byte) (int, error) {
-	if r.pos >= r.failAfter {
-		return 0, io.ErrUnexpectedEOF
-	}
-
-	if r.pos >= len(r.data) {
-		return 0, io.EOF
-	}
-
-	n := copy(buf, r.data[r.pos:])
-	r.pos += n
-
-	if r.pos >= r.failAfter {
-		return n, io.ErrUnexpectedEOF
-	}
-
-	return n, nil
-}
diff --git a/internal/protocol/message_decoder.go b/internal/protocol/message_decoder.go
new file mode 100644
index 0000000..72eea1f
--- /dev/null
+++ b/internal/protocol/message_decoder.go
@@ -0,0 +1,65 @@
+package protocol
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+)
+
+// MessageDecoder handles decoding of protocol message headers
+type MessageDecoder struct{}
+
+// NewMessageDecoder creates a new message decoder
+func NewMessageDecoder() *MessageDecoder {
+	return &MessageDecoder{}
+}
+
+// Decode reads the message header and returns a Message with the payload stream
+func (d *MessageDecoder) Decode(r io.Reader) (*Message, error) {
+	// Read message type (1 byte)
+	var msgType MessageType
+	if err := binary.Read(r, binary.BigEndian, &msgType); err != nil {
+		if err == io.EOF {
+			return nil, err
+		}
+		return nil, fmt.Errorf("failed to read message type: %w", err)
+	}
+
+	// Validate message type (only request types are valid for server)
+	if !isValidRequestType(msgType) {
+		return nil, fmt.Errorf("invalid message type: 0x%02x", msgType)
+	}
+
+	// Read payload length (4 bytes, big-endian)
+	var payloadLength uint32
+	if err := binary.Read(r, binary.BigEndian, &payloadLength); err != nil {
+		return nil, fmt.Errorf("failed to read payload length: %w", err)
+	}
+
+	// Validate payload length
+	if payloadLength > MaxPayloadSize {
+		return nil, fmt.Errorf("payload length %d exceeds maximum %d", payloadLength, MaxPayloadSize)
+	}
+
+	// Create limited reader for the payload
+	var payloadStream io.Reader
+	if payloadLength > 0 {
+		payloadStream = io.LimitReader(r, int64(payloadLength))
+	}
+
+	return &Message{
+		Type:           msgType,
+		PayloadLength:  payloadLength,
+		PayloadStream:  payloadStream,
+	}, nil
+}
+
+// isValidRequestType checks if the message type is a valid request type
+func isValidRequestType(msgType MessageType) bool {
+	switch msgType {
+	case ChallengeRequestType, SolutionRequestType:
+		return true
+	default:
+		return false
+	}
+}
diff --git a/internal/protocol/message_decoder_test.go b/internal/protocol/message_decoder_test.go
new file mode 100644
index 0000000..e4c4680
--- /dev/null
+++ b/internal/protocol/message_decoder_test.go
@@ -0,0 +1,206 @@
+package protocol
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMessageDecoder_Decode_Header(t *testing.T) {
+	decoder := NewMessageDecoder()
+
+	tests := []struct {
+		name           string
+		data           []byte
+		wantType       MessageType
+		wantLength     uint32
+		wantErr        string
+	}{
+		{
+			name:       "challenge request with empty payload",
+			data:       []byte{0x01, 0x00, 0x00, 0x00, 0x00},
+			wantType:   ChallengeRequestType,
+			wantLength: 0,
+		},
+		{
+			name:       "solution request with payload",
+			data:       append([]byte{0x03, 0x00, 0x00, 0x00, 0x05}, []byte("hello")...),
+			wantType:   SolutionRequestType,
+			wantLength: 5,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			buf := bytes.NewBuffer(tt.data)
+			msg, err := decoder.Decode(buf)
+
+			if tt.wantErr != "" {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), tt.wantErr)
+				return
+			}
+
+			require.NoError(t, err)
+			assert.Equal(t, tt.wantType, msg.Type)
+			assert.Equal(t, tt.wantLength, msg.PayloadLength)
+
+			if tt.wantLength > 0 {
+				assert.NotNil(t, msg.PayloadStream)
+			} else {
+				assert.Nil(t, msg.PayloadStream)
+			}
+		})
+	}
+}
+
+func TestMessageDecoder_Decode_Errors(t *testing.T) {
+	decoder := NewMessageDecoder()
+
+	tests := []struct {
+		name    string
+		data    []byte
+		wantErr string
+	}{
+		{
+			name:    "empty data",
+			data:    []byte{},
+			wantErr: "EOF",
+		},
+		{
+			name:    "invalid message type",
+			data:    []byte{0xFF, 0x00, 0x00, 0x00, 0x00},
+			wantErr: "invalid message type",
+		},
+		{
+			name:    "response type not allowed",
+			data:    []byte{0x02, 0x00, 0x00, 0x00, 0x00}, // ChallengeResponseType
+			wantErr: "invalid message type",
+		},
+		{
+			name:    "incomplete header",
+			data:    []byte{0x01, 0x00, 0x00},
+			wantErr: "failed to read payload length",
+		},
+		{
+			name:    "payload too large",
+			data:    append([]byte{0x01}, encodeBigEndianUint32(MaxPayloadSize+1)...),
+			wantErr: "payload length",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			buf := bytes.NewBuffer(tt.data)
+			_, err := decoder.Decode(buf)
+			assert.Error(t, err)
+			assert.Contains(t, err.Error(), tt.wantErr)
+		})
+	}
+}
+
+func TestChallengeRequest_Decode(t *testing.T) {
+	req := &ChallengeRequest{}
+
+	t.Run("always succeeds", func(t *testing.T) {
+		err := req.Decode(nil)
+		assert.NoError(t, err)
+
+		err = req.Decode(bytes.NewReader([]byte("ignored")))
+		assert.NoError(t, err)
+	})
+}
+
+func TestSolutionRequest_Decode(t *testing.T) {
+	req := &SolutionRequest{}
+
+	t.Run("valid solution request", func(t *testing.T) {
+		payload := `{"challenge":{"timestamp":1640995200,"difficulty":4,"resource":"quotes","random":"cmFuZG9tMTIz","hmac":"aG1hY19zaWduYXR1cmU="},"nonce":12345}`
+		reader := bytes.NewReader([]byte(payload))
+
+		err := req.Decode(reader)
+		require.NoError(t, err)
+
+		assert.Equal(t, int64(1640995200), req.Challenge.Timestamp)
+		assert.Equal(t, 4, req.Challenge.Difficulty)
+		assert.Equal(t, "quotes", req.Challenge.Resource)
+		assert.Equal(t, uint64(12345), req.Nonce)
+	})
+
+	t.Run("empty payload should error", func(t *testing.T) {
+		err := req.Decode(nil)
+		assert.Error(t, err)
+	})
+
+	t.Run("invalid JSON should error", func(t *testing.T) {
+		payload := `{invalid json}`
+		reader := bytes.NewReader([]byte(payload))
+
+		err := req.Decode(reader)
+		assert.Error(t, err)
+	})
+
+	t.Run("invalid UTF-8 should error", func(t *testing.T) {
+		payload := []byte{0xFF, 0xFE, 0xFD}
+		reader := bytes.NewReader(payload)
+
+		err := req.Decode(reader)
+		assert.Error(t, err)
+	})
+}
+
+func TestEndToEnd_RequestDecoding(t *testing.T) {
+	decoder := NewMessageDecoder()
+
+	t.Run("challenge request flow", func(t *testing.T) {
+		// Create message data: type=0x01, length=0
+		data := []byte{0x01, 0x00, 0x00, 0x00, 0x00}
+		buf := bytes.NewBuffer(data)
+
+		// Decode header
+		msg, err := decoder.Decode(buf)
+		require.NoError(t, err)
+		assert.Equal(t, ChallengeRequestType, msg.Type)
+
+		// Decode request
+		req := &ChallengeRequest{}
+		err = req.Decode(msg.PayloadStream)
+		require.NoError(t, err)
+	})
+
+	t.Run("solution request flow", func(t *testing.T) {
+		payload := `{"challenge":{"timestamp":1640995200,"difficulty":4,"resource":"quotes","random":"cmFuZG9tMTIz","hmac":"aG1hY19zaWduYXR1cmU="},"nonce":12345}`
+
+		// Create message data: type=0x03, length, payload
+		var buf bytes.Buffer
+		buf.WriteByte(0x03) // SolutionRequestType
+		length := uint32(len(payload))
+		buf.Write(encodeBigEndianUint32(length))
+		buf.WriteString(payload)
+
+		// Decode header
+		msg, err := decoder.Decode(&buf)
+		require.NoError(t, err)
+		assert.Equal(t, SolutionRequestType, msg.Type)
+		assert.Equal(t, length, msg.PayloadLength)
+
+		// Decode request
+		req := &SolutionRequest{}
+		err = req.Decode(msg.PayloadStream)
+		require.NoError(t, err)
+		assert.Equal(t, uint64(12345), req.Nonce)
+	})
+}
+
+// Helper functions for testing
+
+func encodeBigEndianUint32(val uint32) []byte {
+	return []byte{
+		byte(val >> 24),
+		byte(val >> 16),
+		byte(val >> 8),
+		byte(val),
+	}
+}
diff --git a/internal/protocol/types.go b/internal/protocol/types.go
index 4afd17f..bb2a8ee 100644
--- a/internal/protocol/types.go
+++ b/internal/protocol/types.go
@@ -1,64 +1,10 @@
 package protocol
 
-import (
-	"hash-of-wisdom/internal/pow/challenge"
-	"hash-of-wisdom/internal/quotes"
-)
+import "io"
 
-// MessageType represents the type of protocol message
-type MessageType byte
-
-const (
-	ChallengeRequest  MessageType = 0x01
-	ChallengeResponse MessageType = 0x02
-	SolutionRequest   MessageType = 0x03
-	QuoteResponse     MessageType = 0x04
-	ErrorResponse     MessageType = 0x05
-)
-
-// Message represents a protocol message with type and payload
+// Message represents a protocol message with type and payload stream
 type Message struct {
-	Type    MessageType
-	Payload []byte
+	Type           MessageType
+	PayloadLength  uint32
+	PayloadStream  io.Reader
 }
-
-// ChallengeRequestPayload is empty (no payload for challenge requests)
-type ChallengeRequestPayload struct{}
-
-// ChallengeResponsePayload is the direct challenge object (not wrapped)
-type ChallengeResponsePayload challenge.Challenge
-
-// SolutionRequestPayload contains the client's solution attempt
-type SolutionRequestPayload struct {
-	Challenge challenge.Challenge `json:"challenge"`
-	Nonce     uint64              `json:"nonce"`
-}
-
-// QuoteResponsePayload is the direct quote object (not wrapped)
-type QuoteResponsePayload quotes.Quote
-
-// ErrorResponsePayload contains error information
-type ErrorResponsePayload struct {
-	Code       string            `json:"code"`
-	Message    string            `json:"message"`
-	RetryAfter int               `json:"retry_after,omitempty"`
-	Details    map[string]string `json:"details,omitempty"`
-}
-
-// Error codes as defined in protocol specification
-const (
-	ErrMalformedMessage    = "MALFORMED_MESSAGE"
-	ErrInvalidChallenge    = "INVALID_CHALLENGE"
-	ErrInvalidSolution     = "INVALID_SOLUTION"
-	ErrExpiredChallenge    = "EXPIRED_CHALLENGE"
-	ErrRateLimited         = "RATE_LIMITED"
-	ErrServerError         = "SERVER_ERROR"
-	ErrTooManyConnections  = "TOO_MANY_CONNECTIONS"
-	ErrDifficultyTooHigh   = "DIFFICULTY_TOO_HIGH"
-)
-
-// Protocol constants
-const (
-	MaxPayloadSize = 8 * 1024 // 8KB maximum payload size
-	HeaderSize     = 5        // 1 byte type + 4 bytes length
-)
-- 
2.44.1


From d874a2ceb4d1f8cbdaff32f44e98cd8d83c076a5 Mon Sep 17 00:00:00 2001
From: Savely Krendelhoff <alicecandyom@proton.me>
Date: Sat, 23 Aug 2025 12:24:27 +0700
Subject: [PATCH 6/7] [PHASE-5] Implement roundtrip tests

---
 internal/protocol/message_decoder_test.go |  91 ++++++----
 internal/protocol/roundtrip_test.go       | 205 ++++++++++++++++++++++
 2 files changed, 260 insertions(+), 36 deletions(-)
 create mode 100644 internal/protocol/roundtrip_test.go

diff --git a/internal/protocol/message_decoder_test.go b/internal/protocol/message_decoder_test.go
index e4c4680..9b8cf47 100644
--- a/internal/protocol/message_decoder_test.go
+++ b/internal/protocol/message_decoder_test.go
@@ -2,6 +2,7 @@ package protocol
 
 import (
 	"bytes"
+	"io"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -102,53 +103,71 @@ func TestMessageDecoder_Decode_Errors(t *testing.T) {
 }
 
 func TestChallengeRequest_Decode(t *testing.T) {
-	req := &ChallengeRequest{}
+	tests := []struct {
+		name   string
+		stream io.Reader
+	}{
+		{"nil stream", nil},
+		{"non-empty stream", bytes.NewReader([]byte("ignored"))},
+	}
 
-	t.Run("always succeeds", func(t *testing.T) {
-		err := req.Decode(nil)
-		assert.NoError(t, err)
-
-		err = req.Decode(bytes.NewReader([]byte("ignored")))
-		assert.NoError(t, err)
-	})
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := &ChallengeRequest{}
+			err := req.Decode(tt.stream)
+			assert.NoError(t, err)
+		})
+	}
 }
 
 func TestSolutionRequest_Decode(t *testing.T) {
-	req := &SolutionRequest{}
+	tests := []struct {
+		name      string
+		payload   []byte
+		wantErr   bool
+		wantNonce uint64
+	}{
+		{
+			name:      "valid solution request",
+			payload:   []byte(`{"challenge":{"timestamp":1640995200,"difficulty":4,"resource":"quotes","random":"cmFuZG9tMTIz","hmac":"aG1hY19zaWduYXR1cmU="},"nonce":12345}`),
+			wantNonce: 12345,
+		},
+		{
+			name:    "invalid JSON",
+			payload: []byte(`{invalid json}`),
+			wantErr: true,
+		},
+		{
+			name:    "invalid UTF-8",
+			payload: []byte{0xFF, 0xFE, 0xFD},
+			wantErr: true,
+		},
+	}
 
-	t.Run("valid solution request", func(t *testing.T) {
-		payload := `{"challenge":{"timestamp":1640995200,"difficulty":4,"resource":"quotes","random":"cmFuZG9tMTIz","hmac":"aG1hY19zaWduYXR1cmU="},"nonce":12345}`
-		reader := bytes.NewReader([]byte(payload))
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := &SolutionRequest{}
+			var reader io.Reader
+			if tt.payload != nil {
+				reader = bytes.NewReader(tt.payload)
+			}
 
-		err := req.Decode(reader)
-		require.NoError(t, err)
+			err := req.Decode(reader)
 
-		assert.Equal(t, int64(1640995200), req.Challenge.Timestamp)
-		assert.Equal(t, 4, req.Challenge.Difficulty)
-		assert.Equal(t, "quotes", req.Challenge.Resource)
-		assert.Equal(t, uint64(12345), req.Nonce)
-	})
+			if tt.wantErr {
+				assert.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tt.wantNonce, req.Nonce)
+			}
+		})
+	}
 
-	t.Run("empty payload should error", func(t *testing.T) {
+	t.Run("nil stream should error", func(t *testing.T) {
+		req := &SolutionRequest{}
 		err := req.Decode(nil)
 		assert.Error(t, err)
 	})
-
-	t.Run("invalid JSON should error", func(t *testing.T) {
-		payload := `{invalid json}`
-		reader := bytes.NewReader([]byte(payload))
-
-		err := req.Decode(reader)
-		assert.Error(t, err)
-	})
-
-	t.Run("invalid UTF-8 should error", func(t *testing.T) {
-		payload := []byte{0xFF, 0xFE, 0xFD}
-		reader := bytes.NewReader(payload)
-
-		err := req.Decode(reader)
-		assert.Error(t, err)
-	})
 }
 
 func TestEndToEnd_RequestDecoding(t *testing.T) {
diff --git a/internal/protocol/roundtrip_test.go b/internal/protocol/roundtrip_test.go
new file mode 100644
index 0000000..d32f3d8
--- /dev/null
+++ b/internal/protocol/roundtrip_test.go
@@ -0,0 +1,205 @@
+package protocol
+
+import (
+	"bytes"
+	"encoding/binary"
+	"encoding/json"
+	"testing"
+
+	"hash-of-wisdom/internal/pow/challenge"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestChallengeResponse_BinaryFormat(t *testing.T) {
+	testChallenge := &challenge.Challenge{
+		Timestamp:  1640995200,
+		Difficulty: 4,
+		Resource:   "quotes",
+		Random:     []byte("test-random"),
+		HMAC:       []byte("test-hmac"),
+	}
+
+	response := &ChallengeResponse{Challenge: testChallenge}
+	var buf bytes.Buffer
+	err := response.Encode(&buf)
+	require.NoError(t, err)
+
+	// Verify binary format
+	data := buf.Bytes()
+	assert.GreaterOrEqual(t, len(data), 5)
+	assert.Equal(t, byte(ChallengeResponseType), data[0])
+
+	payloadLength := binary.BigEndian.Uint32(data[1:5])
+	assert.Greater(t, payloadLength, uint32(0))
+	assert.Equal(t, len(data)-5, int(payloadLength))
+
+	// Verify payload content
+	payload := string(data[5:])
+	assert.Contains(t, payload, "1640995200")
+	assert.Contains(t, payload, "quotes")
+}
+
+func TestSolutionRequest_RoundTrip(t *testing.T) {
+	decoder := NewMessageDecoder()
+
+	payload := `{"challenge":{"timestamp":1640995200,"difficulty":4,"resource":"quotes","random":"dGVzdC1yYW5kb20=","hmac":"dGVzdC1obWFj"},"nonce":12345}`
+
+	var buf bytes.Buffer
+	buf.WriteByte(byte(SolutionRequestType))
+	binary.Write(&buf, binary.BigEndian, uint32(len(payload)))
+	buf.WriteString(payload)
+
+	// Decode header
+	msg, err := decoder.Decode(&buf)
+	require.NoError(t, err)
+	assert.Equal(t, SolutionRequestType, msg.Type)
+	assert.Equal(t, uint32(len(payload)), msg.PayloadLength)
+
+	// Decode request
+	req := &SolutionRequest{}
+	err = req.Decode(msg.PayloadStream)
+	require.NoError(t, err)
+
+	assert.Equal(t, int64(1640995200), req.Challenge.Timestamp)
+	assert.Equal(t, 4, req.Challenge.Difficulty)
+	assert.Equal(t, "quotes", req.Challenge.Resource)
+	assert.Equal(t, uint64(12345), req.Nonce)
+	assert.Equal(t, []byte("test-random"), req.Challenge.Random)
+	assert.Equal(t, []byte("test-hmac"), req.Challenge.HMAC)
+}
+
+func TestErrorResponse_BinaryFormat(t *testing.T) {
+	errorResp := &ErrorResponse{
+		Code:       "INVALID_SOLUTION",
+		Message:    "Test error message",
+		RetryAfter: 30,
+		Details:    map[string]string{"reason": "test"},
+	}
+
+	var buf bytes.Buffer
+	err := errorResp.Encode(&buf)
+	require.NoError(t, err)
+
+	data := buf.Bytes()
+	assert.Equal(t, byte(ErrorResponseType), data[0])
+
+	length := binary.BigEndian.Uint32(data[1:5])
+	assert.Greater(t, length, uint32(0))
+	assert.LessOrEqual(t, length, uint32(MaxPayloadSize))
+
+	payload := string(data[5:])
+	assert.Contains(t, payload, "INVALID_SOLUTION")
+	assert.Contains(t, payload, "Test error message")
+	assert.Contains(t, payload, "30")
+	assert.Contains(t, payload, "test")
+}
+
+func TestChallengeRequest_EmptyPayload(t *testing.T) {
+	decoder := NewMessageDecoder()
+
+	data := []byte{0x01, 0x00, 0x00, 0x00, 0x00}
+	buf := bytes.NewBuffer(data)
+
+	msg, err := decoder.Decode(buf)
+	require.NoError(t, err)
+	assert.Equal(t, ChallengeRequestType, msg.Type)
+	assert.Equal(t, uint32(0), msg.PayloadLength)
+	assert.Nil(t, msg.PayloadStream)
+
+	req := &ChallengeRequest{}
+	err = req.Decode(msg.PayloadStream)
+	require.NoError(t, err)
+}
+
+func TestMessageDecoder_RejectsResponseTypes(t *testing.T) {
+	decoder := NewMessageDecoder()
+
+	data := []byte{byte(ErrorResponseType), 0x00, 0x00, 0x00, 0x05, 'h', 'e', 'l', 'l', 'o'}
+	buf := bytes.NewBuffer(data)
+
+	_, err := decoder.Decode(buf)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "invalid message type")
+}
+
+func TestPayloadStream_LimitedRead(t *testing.T) {
+	decoder := NewMessageDecoder()
+
+	payload := `{"challenge":{"timestamp":1640995200,"difficulty":4,"resource":"quotes","random":"dGVzdA==","hmac":"dGVzdA=="},"nonce":999}`
+	extraData := "this should not be read"
+
+	var buf bytes.Buffer
+	buf.WriteByte(byte(SolutionRequestType))
+	binary.Write(&buf, binary.BigEndian, uint32(len(payload)))
+	buf.WriteString(payload)
+	buf.WriteString(extraData)
+
+	msg, err := decoder.Decode(&buf)
+	require.NoError(t, err)
+
+	req := &SolutionRequest{}
+	err = req.Decode(msg.PayloadStream)
+	require.NoError(t, err)
+	assert.Equal(t, uint64(999), req.Nonce)
+
+	// Verify extra data wasn't consumed
+	remaining := buf.String()
+	assert.Equal(t, extraData, remaining)
+}
+
+func TestTrueRoundTrip_ServerClientCommunication(t *testing.T) {
+	// Simulate actual server-client communication
+
+	// 1. SERVER: Create and encode a challenge response
+	originalChallenge := &challenge.Challenge{
+		Timestamp:  1640995200,
+		Difficulty: 4,
+		Resource:   "quotes",
+		Random:     []byte("server-random-data"),
+		HMAC:       []byte("server-hmac-signature"),
+	}
+
+	serverResponse := &ChallengeResponse{Challenge: originalChallenge}
+	var networkBuffer bytes.Buffer
+
+	// Server encodes response to "network"
+	err := serverResponse.Encode(&networkBuffer)
+	require.NoError(t, err)
+
+	// 2. NETWORK: Simulate transmission (networkBuffer contains binary data)
+	wireData := networkBuffer.Bytes()
+	assert.Greater(t, len(wireData), 5) // Has header + payload
+
+	// 3. CLIENT: Receives and decodes the binary data
+	// (Client would use a generic decoder, not our server-side MessageDecoder)
+	clientBuf := bytes.NewBuffer(wireData)
+
+	// Client reads header manually
+	var msgType MessageType
+	err = binary.Read(clientBuf, binary.BigEndian, &msgType)
+	require.NoError(t, err)
+	assert.Equal(t, ChallengeResponseType, msgType)
+
+	var payloadLength uint32
+	err = binary.Read(clientBuf, binary.BigEndian, &payloadLength)
+	require.NoError(t, err)
+
+	// Client reads payload
+	payloadBytes := make([]byte, payloadLength)
+	_, err = clientBuf.Read(payloadBytes)
+	require.NoError(t, err)
+
+	// Client deserializes the challenge
+	var receivedChallenge challenge.Challenge
+	err = json.Unmarshal(payloadBytes, &receivedChallenge)
+	require.NoError(t, err)
+
+	// 4. VERIFY: Client received exactly what server sent
+	assert.Equal(t, originalChallenge.Timestamp, receivedChallenge.Timestamp)
+	assert.Equal(t, originalChallenge.Difficulty, receivedChallenge.Difficulty)
+	assert.Equal(t, originalChallenge.Resource, receivedChallenge.Resource)
+	assert.Equal(t, originalChallenge.Random, receivedChallenge.Random)
+	assert.Equal(t, originalChallenge.HMAC, receivedChallenge.HMAC)
+}
-- 
2.44.1


From cf7e6a6f2b4a9e462be98679c3539736d8a5a86a Mon Sep 17 00:00:00 2001
From: Savely Krendelhoff <alicecandyom@proton.me>
Date: Sat, 23 Aug 2025 12:35:38 +0700
Subject: [PATCH 7/7] [PHASE-5] Update implementation plan

---
 docs/IMPLEMENTATION.md | 48 ++++++++++++++++++++++++++----------------
 1 file changed, 30 insertions(+), 18 deletions(-)

diff --git a/docs/IMPLEMENTATION.md b/docs/IMPLEMENTATION.md
index 8673278..3b0e54f 100644
--- a/docs/IMPLEMENTATION.md
+++ b/docs/IMPLEMENTATION.md
@@ -68,23 +68,35 @@
 - [X] Create message validation and bounds checking
 - [X] Write unit tests for protocol components
 
-## Phase 5: Basic Server Architecture
+## Phase 5: Binary Protocol Reworking & Application Layer Integration
+- [X] Refactor protocol codec into streaming MessageDecoder
+- [X] Implement streaming message processing with io.Reader
+- [X] Create request/response encoding and decoding methods
+- [X] Add comprehensive round-trip testing for protocol validation
+- [X] Update application layer to use streaming Message interface
+- [X] Fix application tests for new protocol design
+
+## Phase 6: TCP Server & Connection Management
+- [ ] Implement TCP server with connection handling
+- [ ] Add dual timeout protection:
+  - [ ] Connection timeout (max total connection time)
+  - [ ] Read timeout (max idle time between bytes - slowloris protection)
+- [ ] Implement proper connection lifecycle management
+- [ ] Create protocol state machine for request/response flow
+- [ ] Add graceful connection cleanup and error handling
+- [ ] Implement basic client for testing
+- [ ] Write integration tests for client-server communication
+
+## Phase 7: Basic Server Architecture
 - [ ] Set up structured logging (zerolog/logrus)
 - [ ] Set up metrics collection (prometheus)
 - [ ] Create configuration management
 - [ ] Integrate all components into server architecture
 
-## Phase 6: TCP Protocol & Connection Handling
-- [ ] Implement connection handler with proper error handling
-- [ ] Create protocol state machine
-- [ ] Implement connection lifecycle management
-- [ ] Add connection timeout and lifecycle management
-
-## Phase 7: Server Core & Request Handling
-- [ ] Implement TCP server with connection pooling
-- [ ] Create request router and handler dispatcher
-- [ ] Add connection timeout and lifecycle management
+## Phase 8: Advanced Server Features
+- [ ] Add connection pooling and advanced connection management
 - [ ] Implement graceful shutdown mechanism
+- [ ] Add health check endpoints
 - [ ] Add request/response logging middleware
 - [ ] Create health check endpoints
 - [ ] Write integration tests for server core
@@ -98,7 +110,7 @@
 - [ ] Add monitoring for attack detection
 - [ ] Write tests for protection mechanisms
 
-## Phase 9: Observability & Monitoring
+## Phase 10: Observability & Monitoring
 - [ ] Add structured logging throughout application
 - [ ] Implement metrics for key performance indicators:
   - [ ] Active connections count
@@ -110,7 +122,7 @@
 - [ ] Add error categorization and reporting
 - [ ] Implement health check endpoints
 
-## Phase 10: Configuration & Environment Setup
+## Phase 11: Configuration & Environment Setup
 - [ ] Create configuration structure with validation
 - [ ] Support environment variables and config files
 - [ ] Add configuration for different environments (dev/prod)
@@ -118,7 +130,7 @@
 - [ ] Create deployment configuration templates
 - [ ] Add configuration validation and defaults
 
-## Phase 11: Client Implementation
+## Phase 12: Client Implementation
 - [ ] Create client application structure
 - [ ] Implement PoW solver algorithm
 - [ ] Create client-side protocol implementation
@@ -128,7 +140,7 @@
 - [ ] Add client metrics and logging
 - [ ] Write client unit and integration tests
 
-## Phase 12: Docker & Deployment
+## Phase 13: Docker & Deployment
 - [ ] Create multi-stage Dockerfile for server
 - [ ] Create Dockerfile for client
 - [ ] Create docker-compose.yml for local development
@@ -137,7 +149,7 @@
 - [ ] Add environment-specific configurations
 - [ ] Create deployment documentation
 
-## Phase 13: Testing & Quality Assurance
+## Phase 14: Testing & Quality Assurance
 - [ ] Write comprehensive unit tests (>80% coverage):
   - [ ] PoW algorithm tests
   - [ ] Protocol handler tests
@@ -152,7 +164,7 @@
 - [ ] Add benchmark tests for performance validation
 - [ ] Create stress testing scenarios
 
-## Phase 14: Documentation & Final Polish
+## Phase 15: Documentation & Final Polish
 - [ ] Write comprehensive README with setup instructions
 - [ ] Create API documentation for all interfaces
 - [ ] Add inline code documentation
@@ -161,7 +173,7 @@
 - [ ] Add performance tuning recommendations
 - [ ] Create monitoring and alerting guide
 
-## Phase 15: Production Readiness Checklist
+## Phase 16: Production Readiness Checklist
 - [ ] Security audit of all components
 - [ ] Performance benchmarking and optimization
 - [ ] Memory leak detection and prevention
-- 
2.44.1