package rpc import ( "io" "reflect" "testing" "time" ) func TestIOStream(t *testing.T) { handler := NewNezhaHandler() const testStreamID = "ffffffff-ffff-ffff-ffff-ffffffffffff" handler.CreateStream(testStreamID, 0) userIo, agentIo := newPipeReadWriter(), newPipeReadWriter() defer func() { userIo.Close() agentIo.Close() }() handler.AgentConnected(testStreamID, agentIo) handler.UserConnected(testStreamID, userIo) go handler.StartStream(testStreamID, time.Second*10) cases := [][]byte{ {0, 9, 1, 3, 2, 9, 1, 4, 8}, {3, 1, 3, 5, 2, 9, 5, 13, 53, 23}, make([]byte, 1024), make([]byte, 1024*1024), } t.Run("WriteUserIO", func(t *testing.T) { for i, c := range cases { _, err := userIo.Write(c) if err != nil { t.Fatalf("write to userIo failed at case %d: %v", i, err) } b := make([]byte, len(c)) n, err := agentIo.Read(b) if err != nil { t.Fatalf("read agentIo failed at case %d: %v", i, err) } if !reflect.DeepEqual(c, b[:n]) { t.Fatalf("expected %v, but got %v", c, b[:n]) } } }) t.Run("WriteAgentIO", func(t *testing.T) { for i, c := range cases { _, err := agentIo.Write(c) if err != nil { t.Fatalf("write to agentIo failed at case %d: %v", i, err) } b := make([]byte, len(c)) n, err := userIo.Read(b) if err != nil { t.Fatalf("read userIo failed at case %d: %v", i, err) } if !reflect.DeepEqual(c, b[:n]) { t.Fatalf("Expected %v, but got %v", c, b[:n]) } } }) t.Run("WriteUserIOReadTwice", func(t *testing.T) { data := []byte{1, 2, 3, 4, 5, 6, 7, 8} _, err := agentIo.Write(data) if err != nil { t.Fatalf("write to agentIo failed: %v", err) } b := make([]byte, len(data)/2) n, err := userIo.Read(b) if err != nil { t.Fatalf("read userIo failed: %v", err) } b2 := make([]byte, len(data)-n) _, err = userIo.Read(b2) if err != nil { t.Fatalf("read userIo failed: %v", err) } if !reflect.DeepEqual(data[:len(data)/2], b) { t.Fatalf("expected %v, but got %v", data[:len(data)/2], b) } if !reflect.DeepEqual(data[len(data)/2:], b2) { t.Fatalf("expected %v, but got %v", data[len(data)/2:], b2) } }) } func newPipeReadWriter() io.ReadWriteCloser { r, w := io.Pipe() return struct { io.Reader io.WriteCloser }{r, w} } func TestStreamOwnershipReturnsCreatorUserID(t *testing.T) { h := NewNezhaHandler() h.CreateStream("alice-stream", 100) creator, found := h.StreamOwnership("alice-stream") if !found { t.Fatalf("expected stream to be found after CreateStream") } if creator != 100 { t.Fatalf("expected creator user ID 100, got %d", creator) } } func TestStreamOwnershipReturnsNotFoundForUnknownID(t *testing.T) { h := NewNezhaHandler() if _, found := h.StreamOwnership("nonexistent"); found { t.Fatalf("expected unknown stream id to report not-found") } } func TestStreamOwnershipPreservesPerStreamCreator(t *testing.T) { h := NewNezhaHandler() h.CreateStream("alice-stream", 100) h.CreateStream("bob-stream", 200) aliceCreator, _ := h.StreamOwnership("alice-stream") bobCreator, _ := h.StreamOwnership("bob-stream") if aliceCreator != 100 || bobCreator != 200 { t.Fatalf("expected per-stream creator IDs alice=100 bob=200, got alice=%d bob=%d", aliceCreator, bobCreator) } } func TestIsStreamAuthorizedForUserAllowsCreator(t *testing.T) { h := NewNezhaHandler() h.CreateStream("alice-stream", 100) if !h.IsStreamAuthorizedForUser("alice-stream", 100, false) { t.Fatalf("creator must be authorized to attach to their own stream") } } func TestIsStreamAuthorizedForUserDeniesForeignMember(t *testing.T) { h := NewNezhaHandler() h.CreateStream("alice-stream", 100) if h.IsStreamAuthorizedForUser("alice-stream", 200, false) { t.Fatalf("foreign member must not be authorized — session hijack would be possible") } } func TestIsStreamAuthorizedForUserAllowsAdmin(t *testing.T) { h := NewNezhaHandler() h.CreateStream("alice-stream", 100) if !h.IsStreamAuthorizedForUser("alice-stream", 999, true) { t.Fatalf("admin must be authorized to attach regardless of creator") } } func TestIsStreamAuthorizedForUserDeniesUnknownStream(t *testing.T) { h := NewNezhaHandler() if h.IsStreamAuthorizedForUser("nonexistent", 100, true) { t.Fatalf("unknown stream id must not authorize even admin") } }