249 lines
8.5 KiB
Go
249 lines
8.5 KiB
Go
package xray
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"testing"
|
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/protobuf/proto"
|
|
"google.golang.org/protobuf/reflect/protodesc"
|
|
"google.golang.org/protobuf/reflect/protoreflect"
|
|
"google.golang.org/protobuf/types/descriptorpb"
|
|
"google.golang.org/protobuf/types/dynamicpb"
|
|
)
|
|
|
|
// This file proves the wire-level stats client against a STOCK gRPC server:
|
|
// the server side runs grpc's default protobuf codec with dynamicpb messages
|
|
// built from the xray.app.stats.command descriptor (rebuilt here from the
|
|
// .proto field numbers), so what interops is the real wire contract — field
|
|
// numbers, message shapes, the "proto" content subtype — not just this
|
|
// package's own codec talking to itself.
|
|
|
|
func fieldDesc(name string, num int32, typ descriptorpb.FieldDescriptorProto_Type, label descriptorpb.FieldDescriptorProto_Label, typeName string) *descriptorpb.FieldDescriptorProto {
|
|
f := &descriptorpb.FieldDescriptorProto{
|
|
Name: proto.String(name),
|
|
Number: proto.Int32(num),
|
|
Type: typ.Enum(),
|
|
Label: label.Enum(),
|
|
}
|
|
if typeName != "" {
|
|
f.TypeName = proto.String(typeName)
|
|
}
|
|
return f
|
|
}
|
|
|
|
func statsFileDescriptor(t *testing.T) protoreflect.FileDescriptor {
|
|
t.Helper()
|
|
const (
|
|
strT = descriptorpb.FieldDescriptorProto_TYPE_STRING
|
|
i64T = descriptorpb.FieldDescriptorProto_TYPE_INT64
|
|
boolT = descriptorpb.FieldDescriptorProto_TYPE_BOOL
|
|
msgT = descriptorpb.FieldDescriptorProto_TYPE_MESSAGE
|
|
opt = descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL
|
|
rep = descriptorpb.FieldDescriptorProto_LABEL_REPEATED
|
|
pkg = ".xray.app.stats.command."
|
|
)
|
|
msg := func(name string, fields ...*descriptorpb.FieldDescriptorProto) *descriptorpb.DescriptorProto {
|
|
return &descriptorpb.DescriptorProto{Name: proto.String(name), Field: fields}
|
|
}
|
|
fdp := &descriptorpb.FileDescriptorProto{
|
|
Name: proto.String("stats.proto"),
|
|
Package: proto.String("xray.app.stats.command"),
|
|
Syntax: proto.String("proto3"),
|
|
MessageType: []*descriptorpb.DescriptorProto{
|
|
msg("Stat",
|
|
fieldDesc("name", 1, strT, opt, ""),
|
|
fieldDesc("value", 2, i64T, opt, "")),
|
|
msg("QueryStatsRequest",
|
|
fieldDesc("pattern", 1, strT, opt, ""),
|
|
fieldDesc("reset", 2, boolT, opt, "")),
|
|
msg("QueryStatsResponse",
|
|
fieldDesc("stat", 1, msgT, rep, pkg+"Stat")),
|
|
msg("OnlineIPEntry",
|
|
fieldDesc("ip", 1, strT, opt, ""),
|
|
fieldDesc("last_seen", 2, i64T, opt, "")),
|
|
msg("TrafficUserStat",
|
|
fieldDesc("uplink", 1, i64T, opt, ""),
|
|
fieldDesc("downlink", 2, i64T, opt, "")),
|
|
msg("UserStat",
|
|
fieldDesc("email", 1, strT, opt, ""),
|
|
fieldDesc("ips", 2, msgT, rep, pkg+"OnlineIPEntry"),
|
|
fieldDesc("traffic", 3, msgT, opt, pkg+"TrafficUserStat")),
|
|
msg("GetUsersStatsRequest",
|
|
fieldDesc("include_traffic", 1, boolT, opt, ""),
|
|
fieldDesc("reset", 2, boolT, opt, "")),
|
|
msg("GetUsersStatsResponse",
|
|
fieldDesc("users", 1, msgT, rep, pkg+"UserStat")),
|
|
},
|
|
Service: []*descriptorpb.ServiceDescriptorProto{
|
|
{
|
|
Name: proto.String("StatsService"),
|
|
Method: []*descriptorpb.MethodDescriptorProto{
|
|
{Name: proto.String("QueryStats"), InputType: proto.String(pkg + "QueryStatsRequest"), OutputType: proto.String(pkg + "QueryStatsResponse")},
|
|
{Name: proto.String("GetUsersStats"), InputType: proto.String(pkg + "GetUsersStatsRequest"), OutputType: proto.String(pkg + "GetUsersStatsResponse")},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
fd, err := protodesc.NewFile(fdp, nil)
|
|
if err != nil {
|
|
t.Fatalf("build stats descriptor: %v", err)
|
|
}
|
|
return fd
|
|
}
|
|
|
|
// stockStats is a StatsService served by grpc's stock codec from dynamicpb
|
|
// messages, replaying queued responses per method.
|
|
type stockStats struct {
|
|
fd protoreflect.FileDescriptor
|
|
querySeq []*dynamicpb.Message
|
|
usersSeq []*dynamicpb.Message
|
|
qCalls int
|
|
uCalls int
|
|
}
|
|
|
|
func (s *stockStats) newMsg(name protoreflect.Name) *dynamicpb.Message {
|
|
return dynamicpb.NewMessage(s.fd.Messages().ByName(name))
|
|
}
|
|
|
|
func (s *stockStats) method(name string, reqName protoreflect.Name, pick func() *dynamicpb.Message) grpc.MethodDesc {
|
|
return grpc.MethodDesc{
|
|
MethodName: name,
|
|
Handler: func(_ any, _ context.Context, dec func(any) error, _ grpc.UnaryServerInterceptor) (any, error) {
|
|
req := s.newMsg(reqName)
|
|
if err := dec(req); err != nil {
|
|
return nil, err
|
|
}
|
|
return pick(), nil
|
|
},
|
|
}
|
|
}
|
|
|
|
func startStockStats(t *testing.T, st *stockStats) *XrayAPI {
|
|
t.Helper()
|
|
st.fd = statsFileDescriptor(t)
|
|
lis, err := net.Listen("tcp", "127.0.0.1:0")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
srv := grpc.NewServer() // stock codec: the real interop surface
|
|
srv.RegisterService(&grpc.ServiceDesc{
|
|
ServiceName: "xray.app.stats.command.StatsService",
|
|
HandlerType: (*statsServiceIface)(nil),
|
|
Methods: []grpc.MethodDesc{
|
|
st.method("QueryStats", "QueryStatsRequest", func() *dynamicpb.Message {
|
|
i := st.qCalls
|
|
st.qCalls++
|
|
if i >= len(st.querySeq) {
|
|
i = len(st.querySeq) - 1
|
|
}
|
|
return st.querySeq[i]
|
|
}),
|
|
st.method("GetUsersStats", "GetUsersStatsRequest", func() *dynamicpb.Message {
|
|
i := st.uCalls
|
|
st.uCalls++
|
|
if i >= len(st.usersSeq) {
|
|
i = len(st.usersSeq) - 1
|
|
}
|
|
return st.usersSeq[i]
|
|
}),
|
|
},
|
|
}, st)
|
|
go func() { _ = srv.Serve(lis) }()
|
|
t.Cleanup(srv.Stop)
|
|
|
|
api := &XrayAPI{}
|
|
if err := api.Init(lis.Addr().(*net.TCPAddr).Port); err != nil {
|
|
t.Fatalf("api init: %v", err)
|
|
}
|
|
t.Cleanup(api.Close)
|
|
return api
|
|
}
|
|
|
|
func (s *stockStats) statMsg(name string, value int64) *dynamicpb.Message {
|
|
m := s.newMsg("Stat")
|
|
d := m.Descriptor()
|
|
m.Set(d.Fields().ByName("name"), protoreflect.ValueOfString(name))
|
|
m.Set(d.Fields().ByName("value"), protoreflect.ValueOfInt64(value))
|
|
return m
|
|
}
|
|
|
|
func (s *stockStats) queryResp(stats ...*dynamicpb.Message) *dynamicpb.Message {
|
|
m := s.newMsg("QueryStatsResponse")
|
|
list := m.Mutable(m.Descriptor().Fields().ByName("stat")).List()
|
|
for _, st := range stats {
|
|
list.Append(protoreflect.ValueOfMessage(st.ProtoReflect()))
|
|
}
|
|
return m
|
|
}
|
|
|
|
// TestStatsAgainstStockProtoServer drives GetTraffic and GetOnlineUsers
|
|
// against the stock-codec server: the same deltas and user shapes the typed
|
|
// client produced must come out of the wire-level one.
|
|
func TestStatsAgainstStockProtoServer(t *testing.T) {
|
|
st := &stockStats{}
|
|
st.fd = statsFileDescriptor(t)
|
|
st.querySeq = []*dynamicpb.Message{
|
|
st.queryResp(st.statMsg("user>>>alice>>>traffic>>>uplink", 100)),
|
|
st.queryResp(
|
|
st.statMsg("user>>>alice>>>traffic>>>uplink", 175),
|
|
st.statMsg("inbound>>>in-443>>>traffic>>>downlink", 2048),
|
|
),
|
|
}
|
|
|
|
// One online user with two IPs and a traffic block (field 3 — present in
|
|
// the message, ignored by the panel's OnlineUser shape).
|
|
user := st.newMsg("UserStat")
|
|
ud := user.Descriptor()
|
|
user.Set(ud.Fields().ByName("email"), protoreflect.ValueOfString("alice"))
|
|
ipList := user.Mutable(ud.Fields().ByName("ips")).List()
|
|
for _, e := range []struct {
|
|
ip string
|
|
seen int64
|
|
}{{"203.0.113.7", 111}, {"2001:db8::5", 222}} {
|
|
entry := st.newMsg("OnlineIPEntry")
|
|
ed := entry.Descriptor()
|
|
entry.Set(ed.Fields().ByName("ip"), protoreflect.ValueOfString(e.ip))
|
|
entry.Set(ed.Fields().ByName("last_seen"), protoreflect.ValueOfInt64(e.seen))
|
|
ipList.Append(protoreflect.ValueOfMessage(entry.ProtoReflect()))
|
|
}
|
|
traffic := st.newMsg("TrafficUserStat")
|
|
td := traffic.Descriptor()
|
|
traffic.Set(td.Fields().ByName("uplink"), protoreflect.ValueOfInt64(10))
|
|
traffic.Set(td.Fields().ByName("downlink"), protoreflect.ValueOfInt64(20))
|
|
user.Set(ud.Fields().ByName("traffic"), protoreflect.ValueOfMessage(traffic.ProtoReflect()))
|
|
usersResp := st.newMsg("GetUsersStatsResponse")
|
|
uList := usersResp.Mutable(usersResp.Descriptor().Fields().ByName("users")).List()
|
|
uList.Append(protoreflect.ValueOfMessage(user.ProtoReflect()))
|
|
st.usersSeq = []*dynamicpb.Message{usersResp}
|
|
|
|
api := startStockStats(t, st)
|
|
|
|
if _, _, err := api.GetTraffic(); err != nil {
|
|
t.Fatalf("GetTraffic (baseline): %v", err)
|
|
}
|
|
tags, clients, err := api.GetTraffic()
|
|
if err != nil {
|
|
t.Fatalf("GetTraffic: %v", err)
|
|
}
|
|
if len(clients) != 1 || clients[0].Email != "alice" || clients[0].Up != 75 {
|
|
t.Fatalf("clients = %+v, want alice up=75", clients)
|
|
}
|
|
if len(tags) != 1 || tags[0].Tag != "in-443" || tags[0].Down != 2048 {
|
|
t.Fatalf("tags = %+v, want in-443 down=2048", tags)
|
|
}
|
|
|
|
users, err := api.GetOnlineUsers()
|
|
if err != nil {
|
|
t.Fatalf("GetOnlineUsers: %v", err)
|
|
}
|
|
if len(users) != 1 || users[0].Email != "alice" || len(users[0].IPs) != 2 {
|
|
t.Fatalf("users = %+v, want alice with 2 IPs", users)
|
|
}
|
|
if users[0].IPs[0].IP != "203.0.113.7" || users[0].IPs[0].LastSeen != 111 ||
|
|
users[0].IPs[1].IP != "2001:db8::5" || users[0].IPs[1].LastSeen != 222 {
|
|
t.Fatalf("ips = %+v", users[0].IPs)
|
|
}
|
|
}
|