2025-08-23 08:36:45 +03:00
2 changed files with 260 additions and 36 deletions
--- 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) {
+	tests := []struct {
+		name   string
+		stream io.Reader
+	}{
+		{"nil stream", nil},
+		{"non-empty stream", bytes.NewReader([]byte("ignored"))},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(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")))
+			err := req.Decode(tt.stream)
 			assert.NoError(t, err)
 		})
 	}
+}

 func TestSolutionRequest_Decode(t *testing.T) {
+	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,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(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))
+			var reader io.Reader
+			if tt.payload != nil {
+				reader = bytes.NewReader(tt.payload)
+			}

 			err := req.Decode(reader)
+
+			if tt.wantErr {
+				assert.Error(t, err)
+			} else {
 				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, 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) {
--- a/internal/protocol/roundtrip_test.go
+++ 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)
+}