216 lines
9.7 KiB
Go
216 lines
9.7 KiB
Go
package httpserver
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"git.yuntuo.hk/k2safe/orbit/internal/fault"
|
||
"git.yuntuo.hk/k2safe/orbit/internal/jsonwire"
|
||
"git.yuntuo.hk/k2safe/orbit/internal/testcontract"
|
||
core "git.yuntuo.hk/k2safe/orbit/model"
|
||
"git.yuntuo.hk/k2safe/orbit/platform/auth"
|
||
pulse "git.yuntuo.hk/k2safe/orbit/pulse/model"
|
||
"git.yuntuo.hk/k2safe/orbit/pulse/service"
|
||
"git.yuntuo.hk/k2safe/orbit/pulse/worker"
|
||
"github.com/santhosh-tekuri/jsonschema/v6"
|
||
)
|
||
|
||
// 此存储只接收已通过应用校验的命令，测试不伪造执行成功或当前租约。
|
||
type contractStore struct{ worker.Store }
|
||
|
||
func (contractStore) Connect(context.Context, auth.Token, worker.ConnectInput) (worker.Binding, worker.Welcome, error) {
|
||
return worker.Binding{}, worker.Welcome{}, nil
|
||
}
|
||
func (contractStore) Capacity(context.Context, auth.Token, worker.Binding, worker.CapacityInput) (worker.Welcome, error) {
|
||
return worker.Welcome{}, nil
|
||
}
|
||
func (contractStore) Resume(context.Context, auth.Token, worker.Binding, worker.AttemptRef) (worker.LeaseStatus, error) {
|
||
return worker.LeaseStatus{}, nil
|
||
}
|
||
func (contractStore) Start(context.Context, auth.Token, worker.Binding, core.ID, worker.AttemptRef) (worker.CommandResult, error) {
|
||
return worker.CommandResult{}, nil
|
||
}
|
||
func (contractStore) Renew(context.Context, auth.Token, worker.Binding, core.ID, worker.AttemptRef) (worker.CommandResult, error) {
|
||
return worker.CommandResult{}, nil
|
||
}
|
||
func (contractStore) Finish(context.Context, auth.Token, worker.Binding, core.ID, worker.FinishInput) (worker.CommandResult, error) {
|
||
return worker.CommandResult{}, nil
|
||
}
|
||
|
||
func validateClientContract(raw []byte) error {
|
||
var kind struct {
|
||
Type string `json:"type"`
|
||
}
|
||
if err := json.Unmarshal(raw, &kind); err != nil {
|
||
return err
|
||
}
|
||
f, err := decodeWorkerRequest(raw, kind.Type == "worker.hello")
|
||
if err != nil {
|
||
return err
|
||
}
|
||
svc := worker.New(contractStore{})
|
||
ctx := context.Background()
|
||
token, binding := auth.Token{}, worker.Binding{}
|
||
switch f.Type {
|
||
case "worker.hello":
|
||
var in worker.ConnectInput
|
||
if err := decodeWorkerJSON(f.Body, &in); err != nil {
|
||
return err
|
||
}
|
||
_, _, err = svc.Connect(ctx, token, in)
|
||
case "worker.capacity":
|
||
var in worker.CapacityInput
|
||
if err := decodeWorkerJSON(f.Body, &in); err != nil {
|
||
return err
|
||
}
|
||
_, err = svc.Capacity(ctx, token, binding, in)
|
||
case "worker.resume", "task.start", "task.renew":
|
||
var in worker.AttemptRef
|
||
if err := decodeWorkerJSON(f.Body, &in); err != nil {
|
||
return err
|
||
}
|
||
switch f.Type {
|
||
case "worker.resume":
|
||
_, err = svc.Resume(ctx, token, binding, in)
|
||
case "task.start":
|
||
_, err = svc.Start(ctx, token, binding, f.MessageID, in)
|
||
case "task.renew":
|
||
_, err = svc.Renew(ctx, token, binding, f.MessageID, in)
|
||
}
|
||
case "task.finish":
|
||
var in worker.FinishInput
|
||
if err := decodeWorkerJSON(f.Body, &in); err != nil {
|
||
return err
|
||
}
|
||
_, err = svc.Finish(ctx, token, binding, f.MessageID, in)
|
||
default:
|
||
return fault.Invalid("unsupported worker message type")
|
||
}
|
||
return err
|
||
}
|
||
|
||
func TestWorkerPublishedContract(t *testing.T) {
|
||
id := core.ID(strings.Repeat("1", 32))
|
||
secret := strings.Repeat("A", 43)
|
||
ref := worker.AttemptRef{TaskID: id, AttemptID: id, Epoch: 1, LeaseSecret: secret}
|
||
hello := worker.ConnectInput{InstanceID: id, InstanceSecret: secret, QueueID: id, Capacity: 1, Capabilities: []worker.Capability{{Handler: "text.normalize", HandlerVersion: "v1", PayloadVersion: "v1"}}}
|
||
cases := []struct {
|
||
kind string
|
||
body any
|
||
}{
|
||
{"worker.hello", hello},
|
||
{"worker.capacity", worker.CapacityInput{Sequence: 1, Capacity: 0, Drain: true}},
|
||
{"worker.resume", ref}, {"task.start", ref}, {"task.renew", ref},
|
||
{"task.finish", worker.FinishInput{Ref: ref, Outcome: "SUCCEEDED", Result: json.RawMessage(`{"n":9007199254740993}`)}},
|
||
{"task.finish", worker.FinishInput{Ref: ref, Outcome: "FAILED", ErrorCode: "OBJECT_REJECTED", Result: json.RawMessage(`{}`)}},
|
||
{"task.finish", worker.FinishInput{Ref: ref, Outcome: "UNKNOWN", ErrorCode: "LEASE_LOST", Result: json.RawMessage(`{}`)}},
|
||
{"task.finish", worker.FinishInput{Ref: ref, Outcome: "FAILED", Retryable: true, ErrorCode: "TEMPORARY_REJECTION", Result: json.RawMessage(`{}`)}},
|
||
}
|
||
var valid [][]byte
|
||
for _, tc := range cases {
|
||
body, _ := jsonwire.Marshal(tc.body)
|
||
f := workerFrame{Version: "1", Type: tc.kind, MessageID: id, SessionID: id, Body: body}
|
||
if tc.kind == "worker.hello" {
|
||
f.SessionID = ""
|
||
}
|
||
raw, err := jsonwire.Marshal(f)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
testcontract.WorkerFrame(t, "client", raw)
|
||
if err := validateClientContract(raw); err != nil {
|
||
t.Fatalf("valid %s rejected: %v", tc.kind, err)
|
||
}
|
||
valid = append(valid, raw)
|
||
}
|
||
// 修改报文而非重新构造 Go DTO，覆盖不同语言接入时的字段别名、缺失和类型边界。
|
||
invalid := []struct {
|
||
name string
|
||
raw string
|
||
}{
|
||
{"hello_has_session", strings.Replace(string(valid[0]), `"body":`, `"session_id":"", "body":`, 1)},
|
||
{"response_null", strings.Replace(string(valid[0]), `"body":`, `"error":null, "body":`, 1)},
|
||
{"response_reply", strings.Replace(string(valid[0]), `"body":`, `"reply_to":"", "body":`, 1)},
|
||
{"capability_alias", strings.Replace(string(valid[0]), `"handler":`, `"Handler":`, 1)},
|
||
{"capability_conflicting_alias", strings.Replace(string(valid[0]), `"handler":`, `"Handler":"wrong", "handler":`, 1)},
|
||
{"missing_drain", strings.Replace(string(valid[1]), `,"drain":true`, ``, 1)},
|
||
{"counter_numeric", strings.Replace(string(valid[1]), `"sequence":"1"`, `"sequence":1`, 1)},
|
||
{"counter_leading_zero", strings.Replace(string(valid[1]), `"sequence":"1"`, `"sequence":"01"`, 1)},
|
||
{"counter_overflow", strings.Replace(string(valid[1]), `"sequence":"1"`, `"sequence":"18446744073709551616"`, 1)},
|
||
{"counter_zero", strings.Replace(string(valid[1]), `"sequence":"1"`, `"sequence":"0"`, 1)},
|
||
{"capacity_limit", strings.Replace(string(valid[1]), `"capacity":0`, `"capacity":257`, 1)},
|
||
{"secret_noncanonical", strings.Replace(string(valid[2]), secret, strings.Repeat("A", 42)+"B", 1)},
|
||
{"null_result", strings.Replace(string(valid[5]), `{"n":9007199254740993}`, `null`, 1)},
|
||
{"failure_missing_code", strings.Replace(string(valid[6]), `,"error_code":"OBJECT_REJECTED"`, ``, 1)},
|
||
{"success_with_code", strings.Replace(string(valid[6]), `"FAILED"`, `"SUCCEEDED"`, 1)},
|
||
{"unknown_retry_flag", strings.Replace(string(valid[8]), `"FAILED"`, `"UNKNOWN"`, 1)},
|
||
}
|
||
for _, tc := range invalid {
|
||
t.Run(tc.name, func(t *testing.T) {
|
||
value, err := jsonschema.UnmarshalJSON(strings.NewReader(tc.raw))
|
||
if err != nil {
|
||
t.Fatal("fixture must remain valid JSON")
|
||
}
|
||
if testcontract.WorkerSchema(t, "client").Validate(value) == nil {
|
||
t.Fatal("Schema accepted invalid message")
|
||
}
|
||
if validateClientContract([]byte(tc.raw)) == nil {
|
||
t.Fatal("runtime accepted invalid message")
|
||
}
|
||
})
|
||
}
|
||
// 标准 Schema 无法表达原始字节层重复键和 Unicode 替换，运行时必须独立拒绝。
|
||
for _, raw := range []string{
|
||
strings.Replace(string(valid[1]), `"drain":true`, `"drain":true,"drain":true`, 1),
|
||
strings.Replace(string(valid[5]), `{"n":9007199254740993}`, `{"text":"\ud800"}`, 1),
|
||
strings.Replace(string(valid[5]), `{"n":9007199254740993}`, `{"n":1e999}`, 1),
|
||
} {
|
||
if validateClientContract([]byte(raw)) == nil {
|
||
t.Fatal("runtime lost byte-level guard")
|
||
}
|
||
}
|
||
maximum := strings.Replace(string(valid[1]), `"sequence":"1"`, `"sequence":"18446744073709551615"`, 1)
|
||
testcontract.WorkerFrame(t, "client", []byte(maximum))
|
||
if err := validateClientContract([]byte(maximum)); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
now := time.Date(2026, 9, 17, 8, 0, 0, 0, time.UTC)
|
||
welcome := worker.Welcome{SessionID: id, Epoch: 1, ServerTime: now, ExpiresAt: now.Add(45 * time.Second), LeaseDurationMS: 30000}
|
||
lease := worker.LeaseStatus{Active: true, State: pulse.AttemptRunning, LeaseUntil: now.Add(30 * time.Second), ServerTime: now}
|
||
receipt := worker.Receipt{Operation: "start", TaskState: pulse.TaskRunning, AttemptState: pulse.AttemptRunning, TaskRevision: 3, LeaseUntil: lease.LeaseUntil, Disposition: "APPLIED"}
|
||
task := service.TaskView{ID: id, NamespaceID: id, QueueID: id, State: pulse.TaskLeased, Revision: 2, Handler: "text.normalize", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{}`), AvailableAt: now, DeadlineAt: now.Add(time.Hour), RecoverySafety: pulse.RecoveryReconcileOnly, Retry: pulse.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 30000}, AttemptCount: 1, CreatedAt: now, UpdatedAt: now}
|
||
assignment := worker.Assignment{Task: task, AttemptID: id, Epoch: 1, LeaseSecret: secret, LeaseUntil: lease.LeaseUntil, ServerTime: now}
|
||
for _, tc := range []struct {
|
||
kind string
|
||
body any
|
||
}{
|
||
{"worker.welcome", welcome}, {"ack", welcome}, {"ack", lease},
|
||
{"ack", worker.CommandResult{Receipt: receipt, Lease: lease, Deduplicated: true}}, {"task.assign", assignment},
|
||
} {
|
||
raw, _ := jsonwire.Marshal(tc.body)
|
||
f := workerFrame{Version: "1", Type: tc.kind, MessageID: id, SessionID: id, ReplyTo: id, Body: raw}
|
||
if tc.kind == "task.assign" {
|
||
f.ReplyTo = ""
|
||
}
|
||
frame, _ := jsonwire.Marshal(f)
|
||
testcontract.WorkerFrame(t, "server", frame)
|
||
}
|
||
for _, f := range []workerFrame{
|
||
{Version: "1", Type: "error", MessageID: id, ReplyTo: id, Error: fault.Unauthorized},
|
||
{Version: "1", Type: "error", MessageID: id, SessionID: id, Error: worker.StaleSession},
|
||
{Version: "1", Type: "error", MessageID: id, SessionID: id, ReplyTo: id, Error: worker.LeaseLost},
|
||
} {
|
||
raw, _ := jsonwire.Marshal(f)
|
||
testcontract.WorkerFrame(t, "server", raw)
|
||
}
|
||
invalidResponse, _ := jsonwire.Marshal(workerFrame{Version: "1", Type: "task.assign", MessageID: id, SessionID: id, ReplyTo: id, Body: json.RawMessage(`{}`)})
|
||
value, _ := jsonschema.UnmarshalJSON(strings.NewReader(string(invalidResponse)))
|
||
if testcontract.WorkerSchema(t, "server").Validate(value) == nil {
|
||
t.Fatal("invalid server assignment accepted")
|
||
}
|
||
}
|