iop/apps/agent/internal/localcontrol/protocol_test.go

package localcontrol

import (
	"strings"
	"testing"

	iop "iop/proto/gen/iop"

	"google.golang.org/protobuf/encoding/protowire"
	"google.golang.org/protobuf/proto"
)

func TestProtocolValidationMatrix(t *testing.T) {
	t.Parallel()
	valid := runtimeStatusRequest("message-1")
	unknownOperation := runtimeStatusRequest("message-unknown")
	unknownOperation.Operation = "runtime.future"
	readWithCommand := runtimeStatusRequest("message-command")
	readWithCommand.GetRequest().CommandId = "unexpected"
	mutationWithoutCommand := projectRequest(
		"message-mutation",
		"",
		OperationProjectStart,
		"project-a",
	)

	tests := []struct {
		name  string
		input *iop.AgentLocalEnvelope
		code  string
	}{
		{name: "valid", input: valid},
		{
			name: "unsupported version",
			input: func() *iop.AgentLocalEnvelope {
				value := cloneEnvelope(valid)
				value.ProtocolVersion++
				return value
			}(),
			code: ErrorUnsupportedVersion,
		},
		{
			name: "wrong kind",
			input: func() *iop.AgentLocalEnvelope {
				value := cloneEnvelope(valid)
				value.Kind = iop.AgentLocalKind_AGENT_LOCAL_KIND_EVENT
				return value
			}(),
			code: ErrorMalformedFrame,
		},
		{
			name: "missing message",
			input: func() *iop.AgentLocalEnvelope {
				value := cloneEnvelope(valid)
				value.MessageId = ""
				return value
			}(),
			code: ErrorMalformedFrame,
		},
		{
			name:  "unknown operation",
			input: unknownOperation,
			code:  ErrorUnsupportedOperation,
		},
		{
			name:  "read command id",
			input: readWithCommand,
			code:  ErrorMalformedFrame,
		},
		{
			name:  "mutation missing command id",
			input: mutationWithoutCommand,
			code:  ErrorMalformedFrame,
		},
		{
			name: "replay daemon missing",
			input: func() *iop.AgentLocalEnvelope {
				value := cloneEnvelope(valid)
				cursor := uint64(0)
				value.GetRequest().ReplayAfterSequence = &cursor
				return value
			}(),
			code: ErrorMalformedFrame,
		},
		{
			name: "oversize",
			input: func() *iop.AgentLocalEnvelope {
				value := cloneEnvelope(valid)
				value.MessageId = strings.Repeat("x", MaxFrameBytes)
				return value
			}(),
			code: ErrorMalformedFrame,
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()
			failure := ValidateRequest(test.input)
			if test.code == "" {
				if failure != nil {
					t.Fatalf("ValidateRequest: %v", failure)
				}
				return
			}
			if failure == nil || failure.Code != test.code {
				t.Fatalf("failure = %#v, want code %q", failure, test.code)
			}
		})
	}
}

func TestProtocolRejectsUnknownFields(t *testing.T) {
	t.Parallel()
	payload, err := proto.Marshal(runtimeStatusRequest("message-unknown-field"))
	if err != nil {
		t.Fatalf("Marshal: %v", err)
	}
	payload = protowire.AppendTag(payload, 100, protowire.VarintType)
	payload = protowire.AppendVarint(payload, 1)
	envelope, err := ParseEnvelope(payload)
	if err != nil {
		t.Fatalf("ParseEnvelope: %v", err)
	}
	failure := ValidateRequest(envelope)
	if failure == nil || failure.Code != ErrorMalformedFrame {
		t.Fatalf("failure = %#v, want malformed_frame", failure)
	}
}

func TestProtocolCommandHashUsesImmutableArguments(t *testing.T) {
	t.Parallel()
	first := projectRequest(
		"message-a",
		"command-a",
		OperationProjectStart,
		"project-a",
	)
	second := cloneEnvelope(first)
	second.MessageId = "message-b"
	second.GetRequest().CommandId = "command-b"
	cursor := uint64(9)
	second.GetRequest().ReplayAfterSequence = &cursor
	second.GetRequest().ReplayDaemonId = "daemon-a"

	firstHash, err := immutableCommandHash(first.GetOperation(), first.GetRequest())
	if err != nil {
		t.Fatalf("immutableCommandHash first: %v", err)
	}
	secondHash, err := immutableCommandHash(second.GetOperation(), second.GetRequest())
	if err != nil {
		t.Fatalf("immutableCommandHash second: %v", err)
	}
	if firstHash != secondHash {
		t.Fatalf("hash changed for transport-only fields: %q != %q", firstHash, secondHash)
	}
	second.GetRequest().GetProject().MilestoneId = "milestone-b"
	changedHash, err := immutableCommandHash(second.GetOperation(), second.GetRequest())
	if err != nil {
		t.Fatalf("immutableCommandHash changed: %v", err)
	}
	if changedHash == firstHash {
		t.Fatal("immutable argument change did not change command hash")
	}
}

func TestProtocolErrorsAreBoundedAndSafe(t *testing.T) {
	t.Parallel()
	request := runtimeStatusRequest("message-safe")
	response := errorEnvelope(
		request,
		protocolError(ErrorInternal, strings.Repeat("safe", MaxFrameBytes)),
	)
	if response.GetError() == nil {
		t.Fatal("expected typed error payload")
	}
	if len(response.GetError().GetSafeMessage()) > maxSafeTextBytes {
		t.Fatalf("safe message length = %d", len(response.GetError().GetSafeMessage()))
	}
	if response.GetCorrelationId() != request.GetMessageId() ||
		response.GetError().GetCorrelationId() != request.GetMessageId() {
		t.Fatal("error correlation was not preserved")
	}
}