feat: validate Worker protocol and verify process crash recovery

This commit is contained in:
2026-09-17 15:17:09 +08:00
parent 4d4b2e9f7d
commit 9e3d4be5cc
23 changed files with 1878 additions and 28 deletions

View File

@@ -10,6 +10,7 @@ import (
"testing"
"time"
"git.yuntuo.hk/k2safe/orbit/internal/testcontract"
"git.yuntuo.hk/k2safe/orbit/internal/transport/httpserver"
core "git.yuntuo.hk/k2safe/orbit/model"
pulse "git.yuntuo.hk/k2safe/orbit/pulse/model"
@@ -75,6 +76,7 @@ func (c *testWorkerSocket) send(kind string, body any, id core.ID) core.ID {
if err != nil {
c.t.Fatal(err)
}
testcontract.WorkerFrame(c.t, "client", frame)
ctx, cancel := context.WithTimeout(c.ctx, 5*time.Second)
defer cancel()
if err := c.conn.Write(ctx, websocket.MessageText, frame); err != nil {
@@ -90,6 +92,7 @@ func (c *testWorkerSocket) read() testWorkerFrame {
if err != nil {
c.t.Fatal(err)
}
testcontract.WorkerFrame(c.t, "server", raw)
var frame testWorkerFrame
if kind != websocket.MessageText || json.Unmarshal(raw, &frame) != nil || frame.Version != "1" || frame.MessageID.Validate() != nil {
c.t.Fatal("invalid server frame")
@@ -191,6 +194,10 @@ func TestIntegrationWorkerWebSocket(t *testing.T) {
if err := json.Unmarshal(first.ack(startMessage), &begun); err != nil || !begun.Lease.Active || begun.Lease.State != pulse.AttemptRunning {
t.Fatal("missing durable start")
}
var renewed worker.CommandResult
if err := json.Unmarshal(first.ack(first.send("task.renew", ref, "")), &renewed); err != nil || !renewed.Lease.Active {
t.Fatal("missing renewal confirmation")
}
// 模拟旧节点连接断开;新的 WSS 连接进入不同节点并原子迁移有效尝试。
_ = first.conn.CloseNow()
second := connectTestWorker(t, ctx, serverB, fixture)
@@ -220,13 +227,22 @@ func TestIntegrationWorkerWebSocket(t *testing.T) {
})
t.Run("strict_worker_frame_and_hub_shutdown", func(t *testing.T) {
conn := connectTestWorker(t, ctx, serverB, fixture)
secret, _, err := worker.NewSecret()
if err != nil {
t.Fatal(err)
}
message := conn.send("worker.resume", worker.AttemptRef{TaskID: executionID(t), AttemptID: executionID(t), Epoch: 1, LeaseSecret: secret}, "")
failure := conn.read()
if failure.Type != "error" || failure.ReplyTo != message || failure.Error == nil || failure.Error.Code != "NOT_FOUND" {
t.Fatal("missing structured error frame")
}
raw := []byte(`{"protocol_version":"1","protocol_version":"1","type":"worker.capacity"}`)
if err := conn.conn.Write(ctx, websocket.MessageText, raw); err != nil {
t.Fatal(err)
}
readCtx, stop := context.WithTimeout(ctx, 3*time.Second)
defer stop()
_, _, err := conn.conn.Read(readCtx)
_, _, err = conn.conn.Read(readCtx)
if websocket.CloseStatus(err) != websocket.StatusPolicyViolation {
t.Fatalf("duplicate frame not rejected: %v", err)
}

View File

@@ -0,0 +1,61 @@
// Package testcontract 在回归中校验实际报文与公开 Schema,服务二进制不依赖此包。
package testcontract
import (
"strings"
"sync"
"testing"
"git.yuntuo.hk/k2safe/orbit/protocol/workerv1"
"github.com/santhosh-tekuri/jsonschema/v6"
)
var workerOnce sync.Once
var workerSchemas map[string]*jsonschema.Schema
var workerError error
// WorkerSchema 编译并缓存两个方向的 Schema;开启 date-time 等格式断言。
func WorkerSchema(t testing.TB, direction string) *jsonschema.Schema {
t.Helper()
workerOnce.Do(func() {
compiler := jsonschema.NewCompiler()
compiler.AssertFormat()
doc, err := jsonschema.UnmarshalJSON(strings.NewReader(workerv1.Schema))
if err != nil {
workerError = err
return
}
const resource = "https://git.yuntuo.hk/k2safe/orbit/schema/worker.v1.json"
if err := compiler.AddResource(resource, doc); err != nil {
workerError = err
return
}
workerSchemas = make(map[string]*jsonschema.Schema)
for _, side := range []string{"client", "server"} {
workerSchemas[side], workerError = compiler.Compile(resource + "#/$defs/" + side + "Frame")
if workerError != nil {
return
}
}
})
if workerError != nil {
t.Fatal(workerError)
}
schema := workerSchemas[direction]
if schema == nil {
t.Fatal("unknown Worker protocol direction")
}
return schema
}
// WorkerFrame 验证真实传输的 JSON,不在错误中输出可能包含秘密的报文正文。
func WorkerFrame(t testing.TB, direction string, raw []byte) {
t.Helper()
value, err := jsonschema.UnmarshalJSON(strings.NewReader(string(raw)))
if err != nil {
t.Fatal("Worker frame is not valid JSON")
}
if WorkerSchema(t, direction).Validate(value) != nil {
t.Fatalf("%s Worker frame violates published Schema", direction)
}
}

View File

@@ -214,6 +214,18 @@ func exactFields(raw []byte, t reflect.Type) error {
if t == reflect.TypeOf(time.Time{}) || t == reflect.TypeOf(json.RawMessage{}) {
return nil
}
if t.Kind() == reflect.Slice || t.Kind() == reflect.Array {
var values []json.RawMessage
if err := json.Unmarshal(raw, &values); err != nil {
return err
}
for _, value := range values {
if err := exactFields(value, t.Elem()); err != nil {
return err
}
}
return nil
}
if t.Kind() != reflect.Struct {
return nil
}

View File

@@ -7,6 +7,7 @@ import (
"errors"
"net/http"
"reflect"
"strconv"
"strings"
"sync"
"time"
@@ -142,9 +143,9 @@ func (h *WorkerHub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if err != nil {
return
}
var hello workerFrame
hello, decodeErr := decodeWorkerRequest(raw, true)
var input worker.ConnectInput
if kind != websocket.MessageText || decodeWorkerJSON(raw, &hello) != nil || hello.Version != "1" || hello.Type != "worker.hello" || hello.MessageID.Validate() != nil || hello.SessionID != "" || hello.ReplyTo != "" || hello.Error != nil || decodeWorkerJSON(hello.Body, &input) != nil {
if kind != websocket.MessageText || decodeErr != nil || decodeWorkerJSON(hello.Body, &input) != nil {
_ = conn.Close(websocket.StatusPolicyViolation, "invalid worker hello")
return
}
@@ -177,8 +178,8 @@ func (h *WorkerHub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if err != nil {
return
}
var frame workerFrame
if kind != websocket.MessageText || decodeWorkerJSON(raw, &frame) != nil || frame.Version != "1" || frame.MessageID.Validate() != nil || frame.SessionID != binding.SessionID || frame.ReplyTo != "" || frame.Error != nil {
frame, decodeErr := decodeWorkerRequest(raw, false)
if kind != websocket.MessageText || decodeErr != nil || frame.SessionID != binding.SessionID {
_ = conn.Close(websocket.StatusPolicyViolation, "invalid worker frame")
return
}
@@ -393,6 +394,9 @@ func decodeWorkerJSON(raw []byte, out any) error {
if err := exactFields(raw, reflect.TypeOf(out).Elem()); err != nil {
return fault.Invalid("unknown or invalid worker field")
}
if err := requiredWorkerFields(raw, reflect.TypeOf(out).Elem()); err != nil {
return err
}
decoder := json.NewDecoder(bytes.NewReader(raw))
decoder.DisallowUnknownFields()
if err := decoder.Decode(out); err != nil {
@@ -400,3 +404,89 @@ func decodeWorkerJSON(raw []byte, out any) error {
}
return nil
}
// decodeWorkerRequest 拒绝请求中伪装为响应的字段;空字符串/null 不能绕过方向约束。
func decodeWorkerRequest(raw []byte, hello bool) (workerFrame, error) {
var frame workerFrame
if err := decodeWorkerJSON(raw, &frame); err != nil {
return frame, err
}
if frame.Version != "1" || frame.MessageID.Validate() != nil || len(frame.Body) == 0 {
return frame, fault.Invalid("invalid worker envelope")
}
var fields map[string]json.RawMessage
_ = json.Unmarshal(raw, &fields)
if fields["error"] != nil || fields["reply_to"] != nil {
return frame, fault.Invalid("response fields are not allowed in worker requests")
}
if hello {
if frame.Type != "worker.hello" || fields["session_id"] != nil {
return frame, fault.Invalid("invalid worker hello envelope")
}
} else if frame.SessionID.Validate() != nil {
return frame, fault.Invalid("worker session is required")
}
return frame, nil
}
// requiredWorkerFields 校验必填字段及规范 uint64 字符串,不把省略布尔值等同于显式控制意图。
func requiredWorkerFields(raw []byte, t reflect.Type) error {
if t == reflect.TypeOf(json.RawMessage{}) || t == reflect.TypeOf(time.Time{}) {
return nil
}
if t.Kind() == reflect.Pointer {
return nil
}
if t.Kind() == reflect.Array || t.Kind() == reflect.Slice {
var values []json.RawMessage
if err := json.Unmarshal(raw, &values); err != nil {
return fault.Invalid("invalid worker array")
}
for _, value := range values {
if err := requiredWorkerFields(value, t.Elem()); err != nil {
return err
}
}
return nil
}
if t.Kind() != reflect.Struct {
return nil
}
var fields map[string]json.RawMessage
if json.Unmarshal(raw, &fields) != nil {
return fault.Invalid("invalid worker object")
}
for i := 0; i < t.NumField(); i++ {
field := t.Field(i)
tag := strings.Split(field.Tag.Get("json"), ",")
if tag[0] == "" || tag[0] == "-" {
continue
}
value, found := fields[tag[0]]
optional, numberString := false, false
for _, flag := range tag[1:] {
optional = optional || flag == "omitempty"
numberString = numberString || flag == "string"
}
if !found {
if optional {
continue
}
return fault.Invalid("missing required worker field")
}
if field.Type.Kind() == reflect.Uint64 && numberString {
var text string
if json.Unmarshal(value, &text) != nil {
return fault.Invalid("worker counter must be a decimal string")
}
n, err := strconv.ParseUint(text, 10, 64)
if err != nil || n == 0 || strconv.FormatUint(n, 10) != text {
return fault.Invalid("invalid worker counter")
}
}
if err := requiredWorkerFields(value, field.Type); err != nil {
return err
}
}
return nil
}

View File

@@ -0,0 +1,213 @@
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(`{}`)}},
}
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)},
}
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")
}
}