package notify import ( "bytes" "context" "encoding/binary" "net" "testing" "time" "b612.me/stario" ) type bulkAttachScriptConn struct { readBuf *bytes.Reader writeBuf bytes.Buffer } func newBulkAttachScriptConn(inbound []byte) *bulkAttachScriptConn { return &bulkAttachScriptConn{ readBuf: bytes.NewReader(append([]byte(nil), inbound...)), } } func (c *bulkAttachScriptConn) Read(p []byte) (int, error) { return c.readBuf.Read(p) } func (c *bulkAttachScriptConn) Write(p []byte) (int, error) { return c.writeBuf.Write(p) } func (c *bulkAttachScriptConn) Close() error { return nil } func (c *bulkAttachScriptConn) LocalAddr() net.Addr { return bulkAttachTestAddr("local") } func (c *bulkAttachScriptConn) RemoteAddr() net.Addr { return bulkAttachTestAddr("remote") } func (c *bulkAttachScriptConn) SetDeadline(time.Time) error { return nil } func (c *bulkAttachScriptConn) SetReadDeadline(time.Time) error { return nil } func (c *bulkAttachScriptConn) SetWriteDeadline(time.Time) error { return nil } func (c *bulkAttachScriptConn) writtenBytes() []byte { return append([]byte(nil), c.writeBuf.Bytes()...) } type bulkAttachTestAddr string func (a bulkAttachTestAddr) Network() string { return "tcp" } func (a bulkAttachTestAddr) String() string { return string(a) } func encodeDedicatedRecordForAttachTest(payload []byte) []byte { out := make([]byte, bulkDedicatedRecordHeaderLen+len(payload)) copy(out[:4], bulkDedicatedRecordMagic) binary.BigEndian.PutUint32(out[4:8], uint32(len(payload))) copy(out[bulkDedicatedRecordHeaderLen:], payload) return out } func TestSendDedicatedBulkAttachRequestKeepsCoalescedDedicatedPayloadUnread(t *testing.T) { client := NewClient().(*ClientCommon) UseLegacySecurityClient(client) client.msgID = 100 bulk := newBulkHandle(context.Background(), newBulkRuntime("dedicated-attach-test"), clientFileScope(), BulkOpenRequest{ BulkID: "bulk-attach-test", DataID: 1, Dedicated: true, AttachToken: "attach-token", }, 0, nil, nil, 0, nil, nil, nil, nil, nil) encodedResp, err := client.sequenceEn(bulkAttachResponse{Accepted: true}) if err != nil { t.Fatalf("encode bulkAttachResponse failed: %v", err) } replyFrame, err := encodeDirectSignalFrame(stario.NewQueue(), client.sequenceEn, client.msgEn, client.SecretKey, TransferMsg{ ID: 101, Key: systemBulkAttachKey, Value: encodedResp, Type: MSG_SYS_REPLY, }) if err != nil { t.Fatalf("encode attach reply frame failed: %v", err) } dedicatedPayload := []byte("dedicated-tail-bytes") conn := newBulkAttachScriptConn(append(replyFrame, encodeDedicatedRecordForAttachTest(dedicatedPayload)...)) ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() resp, err := client.sendDedicatedBulkAttachRequest(ctx, conn, bulk) if err != nil { t.Fatalf("sendDedicatedBulkAttachRequest failed: %v", err) } if !resp.Accepted { t.Fatalf("bulk attach response = %+v, want accepted", resp) } parsedReq := stario.NewQueue() var reqMsg TransferMsg if err := parsedReq.ParseMessageOwned(conn.writtenBytes(), "attach-request", func(msgq stario.MsgQueue) error { transfer, err := decodeDirectSignalPayload(client.sequenceDe, client.msgDe, client.SecretKey, msgq.Msg) if err != nil { return err } reqMsg = transfer return nil }); err != nil { t.Fatalf("parse written attach request failed: %v", err) } if reqMsg.Key != systemBulkAttachKey || reqMsg.Type != MSG_SYS_WAIT { t.Fatalf("attach request message mismatch: %+v", reqMsg) } readPayload, err := readBulkDedicatedRecord(conn) if err != nil { t.Fatalf("readBulkDedicatedRecord after attach failed: %v", err) } if !bytes.Equal(readPayload, dedicatedPayload) { t.Fatalf("dedicated payload mismatch: got %q want %q", string(readPayload), string(dedicatedPayload)) } } func TestHandleBulkAttachSystemMessageAcceptedWritesDirectReplyBeforeDedicatedHandoff(t *testing.T) { server := NewServer().(*ServerCommon) UseLegacySecurityServer(server) sidecarLeft, sidecarRight := net.Pipe() defer sidecarRight.Close() current := server.bootstrapAcceptedLogical("dedicated-attach-current", nil, sidecarLeft) if current == nil { t.Fatal("bootstrapAcceptedLogical(current) should return logical") } target := server.bootstrapAcceptedLogical("dedicated-attach-target", nil, nil) if target == nil { t.Fatal("bootstrapAcceptedLogical(target) should return logical") } bulk := newBulkHandle(context.Background(), server.getBulkRuntime(), serverFileScope(target), BulkOpenRequest{ BulkID: "server-dedicated-attach-test", DataID: 7, Dedicated: true, AttachToken: "attach-token", }, 0, target, nil, 0, nil, nil, nil, nil, nil) if err := server.getBulkRuntime().register(serverFileScope(target), bulk); err != nil { t.Fatalf("register bulk runtime failed: %v", err) } reqPayload, err := server.sequenceEn(bulkAttachRequest{ PeerID: target.ID(), BulkID: bulk.ID(), AttachToken: "attach-token", }) if err != nil { t.Fatalf("encode bulkAttachRequest failed: %v", err) } msg := Message{ NetType: NET_SERVER, LogicalConn: current, ClientConn: current.compatClientConn(), TransferMsg: TransferMsg{ ID: 42, Key: systemBulkAttachKey, Value: reqPayload, Type: MSG_SYS_WAIT, }, inboundConn: sidecarLeft, Time: time.Now(), } type attachReplyResult struct { transfer TransferMsg resp bulkAttachResponse err error } replyCh := make(chan attachReplyResult, 1) go func() { _ = sidecarRight.SetReadDeadline(time.Now().Add(time.Second)) replyPayload, err := readDirectSignalFramePayload(sidecarRight) if err != nil { replyCh <- attachReplyResult{err: err} return } transfer, err := decodeDirectSignalPayload(server.sequenceDe, current.msgDeSnapshot(), current.secretKeySnapshot(), replyPayload) if err != nil { replyCh <- attachReplyResult{err: err} return } resp, err := decodeBulkAttachResponse(server.sequenceDe, transfer.Value) replyCh <- attachReplyResult{transfer: transfer, resp: resp, err: err} }() if !server.handleBulkAttachSystemMessage(msg) { t.Fatal("handleBulkAttachSystemMessage should accept dedicated attach message") } var result attachReplyResult select { case result = <-replyCh: case <-time.After(2 * time.Second): t.Fatal("timed out waiting for direct attach reply") } if result.err != nil { t.Fatalf("read direct attach reply failed: %v", result.err) } transfer := result.transfer if transfer.ID != msg.ID || transfer.Key != systemBulkAttachKey || transfer.Type != MSG_SYS_REPLY { t.Fatalf("attach reply mismatch: %+v", transfer) } resp := result.resp if !resp.Accepted || resp.Error != "" { t.Fatalf("bulk attach response = %+v, want accepted", resp) } if got := bulk.dedicatedConnSnapshot(); got != sidecarLeft { t.Fatalf("dedicated conn mismatch: got %v want %v", got, sidecarLeft) } if current.transportAttachedSnapshot() { t.Fatal("attach sidecar logical transport should be detached after handoff") } if got := server.GetLogicalConn(current.ID()); got != nil { t.Fatalf("attach sidecar logical should be removed after handoff, got %+v", got) } }