From 4b5596feb3564713637f7c6d27bb70135201b334 Mon Sep 17 00:00:00 2001 From: Risaro <62798663+Risaro@users.noreply.github.com> Date: Sat, 19 Sep 2026 16:39:10 +0900 Subject: [PATCH] XDRIVE transport: Implement it over a local storage backend (#6745) --- infra/conf/transport_internet.go | 13 + infra/conf/transport_method.go | 42 +- infra/conf/transport_test.go | 34 ++ main/distro/all/all.go | 1 + transport/internet/xdrive/config.pb.go | 213 +++++++++ transport/internet/xdrive/config.proto | 18 + transport/internet/xdrive/conn.go | 139 ++++++ transport/internet/xdrive/local.go | 120 +++++ transport/internet/xdrive/params.go | 69 +++ transport/internet/xdrive/storage.go | 39 ++ transport/internet/xdrive/wal.go | 445 +++++++++++++++++ transport/internet/xdrive/xdrive.go | 308 ++++++++++++ transport/internet/xdrive/xdrive_test.go | 583 +++++++++++++++++++++++ 13 files changed, 2021 insertions(+), 3 deletions(-) create mode 100644 transport/internet/xdrive/config.pb.go create mode 100644 transport/internet/xdrive/config.proto create mode 100644 transport/internet/xdrive/conn.go create mode 100644 transport/internet/xdrive/local.go create mode 100644 transport/internet/xdrive/params.go create mode 100644 transport/internet/xdrive/storage.go create mode 100644 transport/internet/xdrive/wal.go create mode 100644 transport/internet/xdrive/xdrive.go create mode 100644 transport/internet/xdrive/xdrive_test.go diff --git a/infra/conf/transport_internet.go b/infra/conf/transport_internet.go index 7b943a054..56a3b556d 100644 --- a/infra/conf/transport_internet.go +++ b/infra/conf/transport_internet.go @@ -36,6 +36,8 @@ func (p TransportProtocol) Build() (string, error) { return "", errors.PrintRemovedFeatureError("QUIC transport (without web service, etc.)", "XHTTP stream-one H3") case "hysteria": return "hysteria", nil + case "xdrive": + return "xdrive", nil default: return "", errors.New("Config: unknown transport protocol: ", p) } @@ -59,6 +61,7 @@ type StreamConfig struct { WSSettings *WebSocketConfig `json:"wsSettings"` HTTPUPGRADESettings *HttpUpgradeConfig `json:"httpupgradeSettings"` HysteriaSettings *HysteriaConfig `json:"hysteriaSettings"` + XDRIVESettings *XDriveConfig `json:"xdriveSettings"` SocketSettings *SocketConfig `json:"sockopt"` } @@ -192,6 +195,16 @@ func (c *StreamConfig) Build() (*internet.StreamConfig, error) { Settings: serial.ToTypedMessage(hs), }) } + if c.XDRIVESettings != nil { + xs, err := c.XDRIVESettings.Build() + if err != nil { + return nil, errors.New("Failed to build XDRIVE config.").Base(err) + } + config.TransportSettings = append(config.TransportSettings, &internet.TransportConfig{ + ProtocolName: "xdrive", + Settings: serial.ToTypedMessage(xs), + }) + } if c.SocketSettings != nil { ss, err := c.SocketSettings.Build() if err != nil { diff --git a/infra/conf/transport_method.go b/infra/conf/transport_method.go index 93cf9ddad..bf7ea7156 100644 --- a/infra/conf/transport_method.go +++ b/infra/conf/transport_method.go @@ -23,6 +23,7 @@ import ( "github.com/xtls/xray-core/transport/internet/splithttp" "github.com/xtls/xray-core/transport/internet/tcp" "github.com/xtls/xray-core/transport/internet/websocket" + "github.com/xtls/xray-core/transport/internet/xdrive" "google.golang.org/protobuf/proto" ) @@ -796,7 +797,42 @@ func readFileOrString(f string, s []string) ([]byte, error) { } type XDriveConfig struct { - RemoteFolder string `json:"remoteFolder"` - Service string `json:"service"` - Secrets []string `json:"secrets"` + RemoteFolder string `json:"remoteFolder"` + Service string `json:"service"` + Secrets []string `json:"secrets"` + SegmentBytes uint32 `json:"segmentBytes"` + FlushIntervalMs uint32 `json:"flushIntervalMs"` + PollIntervalMs uint32 `json:"pollIntervalMs"` + MaxPollIntervalMs uint32 `json:"maxPollIntervalMs"` + SessionTTLSeconds uint32 `json:"sessionTtlSeconds"` + Concurrency uint32 `json:"concurrency"` + EagerWindowMs uint32 `json:"eagerWindowMs"` + HoleTimeoutMs uint32 `json:"holeTimeoutMs"` +} + +// Build implements Buildable. +func (c *XDriveConfig) Build() (proto.Message, error) { + switch c.Service { + case "local": + case "Google Drive": + if len(c.Secrets) != 3 { + return nil, errors.New("Google Drive needs 3 secrets in order of ClientID, ClientSecret, RefreshToken") + } + default: + return nil, errors.New("unsupported service") + } + config := &xdrive.Config{ + RemoteFolder: c.RemoteFolder, + Service: c.Service, + Secrets: c.Secrets, + SegmentBytes: c.SegmentBytes, + FlushIntervalMs: c.FlushIntervalMs, + PollIntervalMs: c.PollIntervalMs, + MaxPollIntervalMs: c.MaxPollIntervalMs, + SessionTtlSeconds: c.SessionTTLSeconds, + Concurrency: c.Concurrency, + EagerWindowMs: c.EagerWindowMs, + HoleTimeoutMs: c.HoleTimeoutMs, + } + return config, nil } diff --git a/infra/conf/transport_test.go b/infra/conf/transport_test.go index 0814fe6f7..14b29d9f3 100644 --- a/infra/conf/transport_test.go +++ b/infra/conf/transport_test.go @@ -291,3 +291,37 @@ func TestHeaderCustomUDPBuildRejectsExprWithoutArgs(t *testing.T) { t.Fatalf("expected transform arg rejection, got %v", err) } } + +func TestXDriveStreamConfig(t *testing.T) { + config := new(StreamConfig) + if err := json.Unmarshal([]byte(`{ + "method": "xdrive", + "xdriveSettings": { + "remoteFolder": "/tmp/xdrive", + "service": "local" + } + }`), config); err != nil { + t.Fatalf("Unmarshal: %v", err) + } + + built, err := config.Build() + if err != nil { + t.Fatalf("Build: %v", err) + } + if built.ProtocolName != "xdrive" { + t.Fatalf("ProtocolName is %q, want %q", built.ProtocolName, "xdrive") + } + if len(built.TransportSettings) != 1 || built.TransportSettings[0].ProtocolName != "xdrive" { + t.Fatalf("TransportSettings is %v, want a single xdrive entry", built.TransportSettings) + } +} + +func TestXDriveRejectsUnknownService(t *testing.T) { + config := new(XDriveConfig) + if err := json.Unmarshal([]byte(`{"remoteFolder": "/tmp/xdrive", "service": "Dropbox"}`), config); err != nil { + t.Fatalf("Unmarshal: %v", err) + } + if _, err := config.Build(); err == nil { + t.Fatal("Build accepted an unsupported service") + } +} diff --git a/main/distro/all/all.go b/main/distro/all/all.go index accb7b07f..1d183c7f1 100644 --- a/main/distro/all/all.go +++ b/main/distro/all/all.go @@ -60,6 +60,7 @@ import ( _ "github.com/xtls/xray-core/transport/internet/tls" _ "github.com/xtls/xray-core/transport/internet/udp" _ "github.com/xtls/xray-core/transport/internet/websocket" + _ "github.com/xtls/xray-core/transport/internet/xdrive" // Transport headers _ "github.com/xtls/xray-core/transport/internet/headers/http" diff --git a/transport/internet/xdrive/config.pb.go b/transport/internet/xdrive/config.pb.go new file mode 100644 index 000000000..6eec3ecd5 --- /dev/null +++ b/transport/internet/xdrive/config.pb.go @@ -0,0 +1,213 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.36.11 +// protoc v6.33.5 +// source: transport/internet/xdrive/config.proto + +package xdrive + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" + unsafe "unsafe" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type Config struct { + state protoimpl.MessageState `protogen:"open.v1"` + RemoteFolder string `protobuf:"bytes,1,opt,name=remote_folder,json=remoteFolder,proto3" json:"remote_folder,omitempty"` + Service string `protobuf:"bytes,2,opt,name=service,proto3" json:"service,omitempty"` + Secrets []string `protobuf:"bytes,3,rep,name=secrets,proto3" json:"secrets,omitempty"` + SegmentBytes uint32 `protobuf:"varint,4,opt,name=segment_bytes,json=segmentBytes,proto3" json:"segment_bytes,omitempty"` + FlushIntervalMs uint32 `protobuf:"varint,5,opt,name=flush_interval_ms,json=flushIntervalMs,proto3" json:"flush_interval_ms,omitempty"` + PollIntervalMs uint32 `protobuf:"varint,6,opt,name=poll_interval_ms,json=pollIntervalMs,proto3" json:"poll_interval_ms,omitempty"` + MaxPollIntervalMs uint32 `protobuf:"varint,7,opt,name=max_poll_interval_ms,json=maxPollIntervalMs,proto3" json:"max_poll_interval_ms,omitempty"` + SessionTtlSeconds uint32 `protobuf:"varint,8,opt,name=session_ttl_seconds,json=sessionTtlSeconds,proto3" json:"session_ttl_seconds,omitempty"` + Concurrency uint32 `protobuf:"varint,9,opt,name=concurrency,proto3" json:"concurrency,omitempty"` + EagerWindowMs uint32 `protobuf:"varint,10,opt,name=eager_window_ms,json=eagerWindowMs,proto3" json:"eager_window_ms,omitempty"` + HoleTimeoutMs uint32 `protobuf:"varint,11,opt,name=hole_timeout_ms,json=holeTimeoutMs,proto3" json:"hole_timeout_ms,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Config) Reset() { + *x = Config{} + mi := &file_transport_internet_xdrive_config_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Config) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Config) ProtoMessage() {} + +func (x *Config) ProtoReflect() protoreflect.Message { + mi := &file_transport_internet_xdrive_config_proto_msgTypes[0] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use Config.ProtoReflect.Descriptor instead. +func (*Config) Descriptor() ([]byte, []int) { + return file_transport_internet_xdrive_config_proto_rawDescGZIP(), []int{0} +} + +func (x *Config) GetRemoteFolder() string { + if x != nil { + return x.RemoteFolder + } + return "" +} + +func (x *Config) GetService() string { + if x != nil { + return x.Service + } + return "" +} + +func (x *Config) GetSecrets() []string { + if x != nil { + return x.Secrets + } + return nil +} + +func (x *Config) GetSegmentBytes() uint32 { + if x != nil { + return x.SegmentBytes + } + return 0 +} + +func (x *Config) GetFlushIntervalMs() uint32 { + if x != nil { + return x.FlushIntervalMs + } + return 0 +} + +func (x *Config) GetPollIntervalMs() uint32 { + if x != nil { + return x.PollIntervalMs + } + return 0 +} + +func (x *Config) GetMaxPollIntervalMs() uint32 { + if x != nil { + return x.MaxPollIntervalMs + } + return 0 +} + +func (x *Config) GetSessionTtlSeconds() uint32 { + if x != nil { + return x.SessionTtlSeconds + } + return 0 +} + +func (x *Config) GetConcurrency() uint32 { + if x != nil { + return x.Concurrency + } + return 0 +} + +func (x *Config) GetEagerWindowMs() uint32 { + if x != nil { + return x.EagerWindowMs + } + return 0 +} + +func (x *Config) GetHoleTimeoutMs() uint32 { + if x != nil { + return x.HoleTimeoutMs + } + return 0 +} + +var File_transport_internet_xdrive_config_proto protoreflect.FileDescriptor + +const file_transport_internet_xdrive_config_proto_rawDesc = "" + + "\n" + + "&transport/internet/xdrive/config.proto\x12\x1exray.transport.internet.xdrive\"\xaf\x03\n" + + "\x06Config\x12#\n" + + "\rremote_folder\x18\x01 \x01(\tR\fremoteFolder\x12\x18\n" + + "\aservice\x18\x02 \x01(\tR\aservice\x12\x18\n" + + "\asecrets\x18\x03 \x03(\tR\asecrets\x12#\n" + + "\rsegment_bytes\x18\x04 \x01(\rR\fsegmentBytes\x12*\n" + + "\x11flush_interval_ms\x18\x05 \x01(\rR\x0fflushIntervalMs\x12(\n" + + "\x10poll_interval_ms\x18\x06 \x01(\rR\x0epollIntervalMs\x12/\n" + + "\x14max_poll_interval_ms\x18\a \x01(\rR\x11maxPollIntervalMs\x12.\n" + + "\x13session_ttl_seconds\x18\b \x01(\rR\x11sessionTtlSeconds\x12 \n" + + "\vconcurrency\x18\t \x01(\rR\vconcurrency\x12&\n" + + "\x0feager_window_ms\x18\n" + + " \x01(\rR\reagerWindowMs\x12&\n" + + "\x0fhole_timeout_ms\x18\v \x01(\rR\rholeTimeoutMsB5Z3github.com/xtls/xray-core/transport/internet/xdriveb\x06proto3" + +var ( + file_transport_internet_xdrive_config_proto_rawDescOnce sync.Once + file_transport_internet_xdrive_config_proto_rawDescData []byte +) + +func file_transport_internet_xdrive_config_proto_rawDescGZIP() []byte { + file_transport_internet_xdrive_config_proto_rawDescOnce.Do(func() { + file_transport_internet_xdrive_config_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_transport_internet_xdrive_config_proto_rawDesc), len(file_transport_internet_xdrive_config_proto_rawDesc))) + }) + return file_transport_internet_xdrive_config_proto_rawDescData +} + +var file_transport_internet_xdrive_config_proto_msgTypes = make([]protoimpl.MessageInfo, 1) +var file_transport_internet_xdrive_config_proto_goTypes = []any{ + (*Config)(nil), // 0: xray.transport.internet.xdrive.Config +} +var file_transport_internet_xdrive_config_proto_depIdxs = []int32{ + 0, // [0:0] is the sub-list for method output_type + 0, // [0:0] is the sub-list for method input_type + 0, // [0:0] is the sub-list for extension type_name + 0, // [0:0] is the sub-list for extension extendee + 0, // [0:0] is the sub-list for field type_name +} + +func init() { file_transport_internet_xdrive_config_proto_init() } +func file_transport_internet_xdrive_config_proto_init() { + if File_transport_internet_xdrive_config_proto != nil { + return + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_transport_internet_xdrive_config_proto_rawDesc), len(file_transport_internet_xdrive_config_proto_rawDesc)), + NumEnums: 0, + NumMessages: 1, + NumExtensions: 0, + NumServices: 0, + }, + GoTypes: file_transport_internet_xdrive_config_proto_goTypes, + DependencyIndexes: file_transport_internet_xdrive_config_proto_depIdxs, + MessageInfos: file_transport_internet_xdrive_config_proto_msgTypes, + }.Build() + File_transport_internet_xdrive_config_proto = out.File + file_transport_internet_xdrive_config_proto_goTypes = nil + file_transport_internet_xdrive_config_proto_depIdxs = nil +} diff --git a/transport/internet/xdrive/config.proto b/transport/internet/xdrive/config.proto new file mode 100644 index 000000000..09133d7c1 --- /dev/null +++ b/transport/internet/xdrive/config.proto @@ -0,0 +1,18 @@ +syntax = "proto3"; + +package xray.transport.internet.xdrive; +option go_package = "github.com/xtls/xray-core/transport/internet/xdrive"; + +message Config { + string remote_folder = 1; + string service = 2; + repeated string secrets = 3; + uint32 segment_bytes = 4; + uint32 flush_interval_ms = 5; + uint32 poll_interval_ms = 6; + uint32 max_poll_interval_ms = 7; + uint32 session_ttl_seconds = 8; + uint32 concurrency = 9; + uint32 eager_window_ms = 10; + uint32 hole_timeout_ms = 11; +} diff --git a/transport/internet/xdrive/conn.go b/transport/internet/xdrive/conn.go new file mode 100644 index 000000000..50f871718 --- /dev/null +++ b/transport/internet/xdrive/conn.go @@ -0,0 +1,139 @@ +package xdrive + +import ( + "context" + "os" + "sync" + "time" + + "github.com/xtls/xray-core/common/net" +) + +var placeholderAddr = &net.TCPAddr{IP: net.IP{127, 0, 0, 1}, Port: 0} + +type Conn struct { + cancel context.CancelFunc + writer *walWriter + reader *walReader + onClose func() + + readBuf []byte + + deadlineMu sync.Mutex + readDeadline time.Time + writeDeadline time.Time + + closeOnce sync.Once + closeErr error +} + +func newConn(ctx context.Context, storage Storage, writePrefix, readPrefix string, p params, onClose func()) *Conn { + ctx, cancel := context.WithCancel(ctx) + return &Conn{ + cancel: cancel, + writer: newWALWriter(ctx, storage, writePrefix, p), + reader: newWALReader(ctx, storage, readPrefix, p), + onClose: onClose, + } +} + +func (c *Conn) Read(b []byte) (int, error) { + if len(c.readBuf) == 0 { + data, err := c.receive() + if err != nil { + return 0, err + } + c.readBuf = data + } + n := copy(b, c.readBuf) + c.readBuf = c.readBuf[n:] + return n, nil +} + +func (c *Conn) receive() ([]byte, error) { + deadline := c.getDeadline(true) + if deadline.IsZero() { + data, ok := <-c.reader.ch + if !ok { + return nil, c.reader.Err() + } + return data, nil + } + + if !time.Now().Before(deadline) { + return nil, os.ErrDeadlineExceeded + } + timer := time.NewTimer(time.Until(deadline)) + defer timer.Stop() + + select { + case data, ok := <-c.reader.ch: + if !ok { + return nil, c.reader.Err() + } + return data, nil + case <-timer.C: + return nil, os.ErrDeadlineExceeded + } +} + +func (c *Conn) Write(b []byte) (int, error) { + if deadline := c.getDeadline(false); !deadline.IsZero() && !time.Now().Before(deadline) { + return 0, os.ErrDeadlineExceeded + } + n, err := c.writer.Write(b) + if err == nil { + c.reader.Wake() + } + return n, err +} + +func (c *Conn) Close() error { + c.closeOnce.Do(func() { + c.closeErr = c.writer.Close() + c.cancel() + if c.onClose != nil { + c.onClose() + } + }) + return c.closeErr +} + +func (c *Conn) LocalAddr() net.Addr { + return placeholderAddr +} + +func (c *Conn) RemoteAddr() net.Addr { + return placeholderAddr +} + +func (c *Conn) getDeadline(read bool) time.Time { + c.deadlineMu.Lock() + defer c.deadlineMu.Unlock() + if read { + return c.readDeadline + } + return c.writeDeadline +} + +func (c *Conn) SetDeadline(t time.Time) error { + c.deadlineMu.Lock() + defer c.deadlineMu.Unlock() + c.readDeadline = t + c.writeDeadline = t + return nil +} + +func (c *Conn) SetReadDeadline(t time.Time) error { + c.deadlineMu.Lock() + defer c.deadlineMu.Unlock() + c.readDeadline = t + return nil +} + +func (c *Conn) SetWriteDeadline(t time.Time) error { + c.deadlineMu.Lock() + defer c.deadlineMu.Unlock() + c.writeDeadline = t + return nil +} diff --git a/transport/internet/xdrive/local.go b/transport/internet/xdrive/local.go new file mode 100644 index 000000000..28757a10e --- /dev/null +++ b/transport/internet/xdrive/local.go @@ -0,0 +1,120 @@ +package xdrive + +import ( + "context" + "os" + "path" + "path/filepath" + "strings" + + "github.com/xtls/xray-core/common/errors" +) + +const tempPrefix = ".xdrive-tmp-" + +type localStorage struct { + root string +} + +func newLocalStorage(root string) (*localStorage, error) { + if root == "" { + return nil, errors.New(`empty "remoteFolder"`) + } + if err := os.MkdirAll(root, 0o700); err != nil { + return nil, errors.New("failed to create remote folder").Base(err) + } + return &localStorage{root: root}, nil +} + +func (s *localStorage) resolve(name string) (string, error) { + clean := path.Clean("/" + name) + if clean == "/" { + return "", errors.New("invalid object name: ", name) + } + if strings.HasPrefix(path.Base(clean), tempPrefix) { + return "", errors.New("reserved object name: ", name) + } + return filepath.Join(s.root, filepath.FromSlash(clean[1:])), nil +} + +func (s *localStorage) Put(ctx context.Context, name string, data []byte) error { + full, err := s.resolve(name) + if err != nil { + return err + } + dir := filepath.Dir(full) + if err := os.MkdirAll(dir, 0o700); err != nil { + return errors.New("failed to create folder ", dir).Base(err) + } + + tmp, err := os.CreateTemp(dir, tempPrefix+"*") + if err != nil { + return errors.New("failed to create temp file in ", dir).Base(err) + } + tmpName := tmp.Name() + defer os.Remove(tmpName) + + if _, err := tmp.Write(data); err != nil { + tmp.Close() + return errors.New("failed to write ", name).Base(err) + } + if err := tmp.Close(); err != nil { + return errors.New("failed to close ", name).Base(err) + } + if err := os.Rename(tmpName, full); err != nil { + return errors.New("failed to commit ", name).Base(err) + } + return nil +} + +func (s *localStorage) Get(ctx context.Context, name string) ([]byte, error) { + full, err := s.resolve(name) + if err != nil { + return nil, err + } + data, err := os.ReadFile(full) + if err != nil { + if os.IsNotExist(err) { + return nil, errNotFound + } + return nil, errors.New("failed to read ", name).Base(err) + } + return data, nil +} + +func (s *localStorage) Delete(ctx context.Context, name string) error { + full, err := s.resolve(name) + if err != nil { + return err + } + if err := os.RemoveAll(full); err != nil { + return errors.New("failed to delete ", name).Base(err) + } + return nil +} + +func (s *localStorage) List(ctx context.Context, prefix string) ([]Entry, error) { + full, err := s.resolve(prefix) + if err != nil { + return nil, err + } + entries, err := os.ReadDir(full) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, errors.New("failed to list ", prefix).Base(err) + } + found := make([]Entry, 0, len(entries)) + for _, entry := range entries { + if strings.HasPrefix(entry.Name(), tempPrefix) { + continue + } + found = append(found, Entry{Name: entry.Name()}) + } + return found, nil +} + +func (s *localStorage) Close() error { + return nil +} diff --git a/transport/internet/xdrive/params.go b/transport/internet/xdrive/params.go new file mode 100644 index 000000000..01f80a31a --- /dev/null +++ b/transport/internet/xdrive/params.go @@ -0,0 +1,69 @@ +package xdrive + +import "time" + +const ( + defaultSegmentBytes = 512 * 1024 + defaultFlushInterval = 20 * time.Millisecond + defaultMinPollInterval = 50 * time.Millisecond + defaultMaxPollInterval = 500 * time.Millisecond + defaultEagerWindow = 2 * time.Second + defaultHoleTimeout = 30 * time.Second + defaultSessionTTL = 5 * time.Minute + defaultConcurrency = 8 + + maxSegmentBytes = 16 * 1024 * 1024 + maxConcurrency = 64 +) + +type params struct { + segmentBytes int + flushInterval time.Duration + minPollInterval time.Duration + maxPollInterval time.Duration + eagerWindow time.Duration + holeTimeout time.Duration + sessionTTL time.Duration + concurrency int +} + +func millis(value uint32, fallback time.Duration) time.Duration { + if value == 0 { + return fallback + } + return time.Duration(value) * time.Millisecond +} + +func seconds(value uint32, fallback time.Duration) time.Duration { + if value == 0 { + return fallback + } + return time.Duration(value) * time.Second +} + +func capped(value uint32, fallback, limit int) int { + if value == 0 { + return fallback + } + if int(value) > limit { + return limit + } + return int(value) +} + +func paramsFromConfig(c *Config) params { + p := params{ + segmentBytes: capped(c.SegmentBytes, defaultSegmentBytes, maxSegmentBytes), + flushInterval: millis(c.FlushIntervalMs, defaultFlushInterval), + minPollInterval: millis(c.PollIntervalMs, defaultMinPollInterval), + maxPollInterval: millis(c.MaxPollIntervalMs, defaultMaxPollInterval), + eagerWindow: millis(c.EagerWindowMs, defaultEagerWindow), + holeTimeout: millis(c.HoleTimeoutMs, defaultHoleTimeout), + sessionTTL: seconds(c.SessionTtlSeconds, defaultSessionTTL), + concurrency: capped(c.Concurrency, defaultConcurrency, maxConcurrency), + } + if p.maxPollInterval < p.minPollInterval { + p.maxPollInterval = p.minPollInterval + } + return p +} diff --git a/transport/internet/xdrive/storage.go b/transport/internet/xdrive/storage.go new file mode 100644 index 000000000..08ea67786 --- /dev/null +++ b/transport/internet/xdrive/storage.go @@ -0,0 +1,39 @@ +package xdrive + +import ( + "context" + + "github.com/xtls/xray-core/common/errors" + "github.com/xtls/xray-core/transport/internet" +) + +var errNotFound = errors.New("object not found") + +type Entry struct { + Name string + Inline []byte +} + +type Storage interface { + Put(ctx context.Context, name string, data []byte) error + Get(ctx context.Context, name string) ([]byte, error) + Delete(ctx context.Context, name string) error + List(ctx context.Context, prefix string) ([]Entry, error) + Close() error +} + +func newStorage(streamSettings *internet.MemoryStreamConfig) (Storage, error) { + config, err := streamConfig(streamSettings) + if err != nil { + return nil, err + } + + switch config.Service { + case "local": + return newLocalStorage(config.RemoteFolder) + case "Google Drive": + return nil, errors.New(`service "Google Drive" is not implemented yet`) + default: + return nil, errors.New("unsupported service: ", config.Service) + } +} diff --git a/transport/internet/xdrive/wal.go b/transport/internet/xdrive/wal.go new file mode 100644 index 000000000..62e1876f2 --- /dev/null +++ b/transport/internet/xdrive/wal.go @@ -0,0 +1,445 @@ +package xdrive + +import ( + "context" + "fmt" + "io" + "strconv" + "strings" + "sync" + "time" + + "github.com/xtls/xray-core/common/errors" +) + +const ( + maxCoalescedTicks = 8 + + segSuffix = ".seg" + endSuffix = ".end" + errSuffix = ".err" +) + +func objectName(prefix string, seq int64, suffix string) string { + return fmt.Sprintf("%s/%09d%s", prefix, seq, suffix) +} + +func parseEntry(name string) (int64, bool) { + dot := strings.LastIndexByte(name, '.') + if dot < 0 { + return 0, false + } + switch name[dot:] { + case segSuffix, endSuffix, errSuffix: + default: + return 0, false + } + seq, err := strconv.ParseInt(name[:dot], 10, 64) + if err != nil || seq < 0 { + return 0, false + } + return seq, true +} + +type walWriter struct { + ctx context.Context + storage Storage + prefix string + params + sem chan struct{} + wg sync.WaitGroup + + mu sync.Mutex + buf []byte + seq int64 + lastSize int + held int + closed bool + err error +} + +func newWALWriter(ctx context.Context, storage Storage, prefix string, p params) *walWriter { + w := &walWriter{ + ctx: ctx, + storage: storage, + prefix: prefix, + params: p, + sem: make(chan struct{}, p.concurrency), + } + go w.flushLoop() + return w +} + +func (w *walWriter) Write(p []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + + if w.err != nil { + return 0, w.err + } + if w.closed { + return 0, io.ErrClosedPipe + } + + w.buf = append(w.buf, p...) + for len(w.buf) >= w.segmentBytes { + if err := w.flushLocked(); err != nil { + return 0, err + } + } + return len(p), nil +} + +func (w *walWriter) flushLoop() { + ticker := time.NewTicker(w.flushInterval) + defer ticker.Stop() + + for { + select { + case <-w.ctx.Done(): + return + case <-ticker.C: + w.mu.Lock() + if !w.closed && w.err == nil && len(w.buf) > 0 && w.readyToFlush() { + w.flushLocked() + } + done := w.closed || w.err != nil + w.mu.Unlock() + if done { + return + } + } + } +} + +func (w *walWriter) readyToFlush() bool { + grew := len(w.buf) > w.lastSize + w.lastSize = len(w.buf) + + if grew && w.held < maxCoalescedTicks { + w.held++ + return false + } + w.held = 0 + return true +} + +func (w *walWriter) flushLocked() error { + if len(w.buf) == 0 { + return nil + } + if w.err != nil { + return w.err + } + + n := len(w.buf) + if n > w.segmentBytes { + n = w.segmentBytes + } + + chunk := make([]byte, n) + copy(chunk, w.buf[:n]) + seq := w.seq + w.seq++ + + if n == len(w.buf) { + w.buf = w.buf[:0] + } else { + w.buf = append(w.buf[:0], w.buf[n:]...) + } + w.lastSize = len(w.buf) + w.held = 0 + + w.wg.Add(1) + go w.upload(seq, chunk) + return nil +} + +func (w *walWriter) upload(seq int64, chunk []byte) { + defer w.wg.Done() + + select { + case w.sem <- struct{}{}: + case <-w.ctx.Done(): + return + } + defer func() { <-w.sem }() + + if err := w.storage.Put(w.ctx, objectName(w.prefix, seq, segSuffix), chunk); err != nil { + w.mu.Lock() + if w.err == nil { + w.err = errors.New("failed to store segment").Base(err) + } + w.mu.Unlock() + w.storage.Put(w.ctx, objectName(w.prefix, seq, errSuffix), nil) + } +} + +func (w *walWriter) Close() error { + w.mu.Lock() + if w.closed { + w.mu.Unlock() + return nil + } + w.closed = true + for len(w.buf) > 0 && w.err == nil { + w.flushLocked() + } + w.mu.Unlock() + + w.wg.Wait() + + w.mu.Lock() + err, seq := w.err, w.seq + w.mu.Unlock() + + if err != nil { + return err + } + return w.storage.Put(w.ctx, objectName(w.prefix, seq, endSuffix), nil) +} + +type walReader struct { + ctx context.Context + storage Storage + prefix string + params + seq int64 + + ch chan []byte + discards chan string + wake chan struct{} + holeSince time.Time + + errMu sync.Mutex + err error +} + +func newWALReader(ctx context.Context, storage Storage, prefix string, p params) *walReader { + r := &walReader{ + ctx: ctx, + storage: storage, + prefix: prefix, + params: p, + ch: make(chan []byte, p.concurrency), + discards: make(chan string, 4*p.concurrency), + wake: make(chan struct{}, 1), + } + go r.run() + go r.discardLoop() + return r +} + +func (r *walReader) Wake() { + select { + case r.wake <- struct{}{}: + default: + } +} + +func (r *walReader) run() { + defer close(r.ch) + + delay := r.minPollInterval + active := time.Now() + for { + polled := time.Now() + advanced, eof, err := r.poll() + if err != nil { + r.setErr(err) + return + } + if eof { + return + } + + switch { + case advanced: + active = time.Now() + delay = r.minPollInterval + case time.Since(active) < r.eagerWindow: + delay = r.minPollInterval + default: + delay *= 2 + if delay > r.maxPollInterval { + delay = r.maxPollInterval + } + } + + timer := time.NewTimer(delay) + select { + case <-r.ctx.Done(): + timer.Stop() + return + case <-timer.C: + case <-r.wake: + timer.Stop() + active = time.Now() + delay = r.minPollInterval + if rest := r.minPollInterval - time.Since(polled); rest > 0 { + select { + case <-r.ctx.Done(): + return + case <-time.After(rest): + } + } + } + } +} + +func (r *walReader) poll() (advanced, eof bool, err error) { + listed, err := r.storage.List(r.ctx, r.prefix) + if err != nil { + return false, false, err + } + if len(listed) == 0 { + return false, false, nil + } + + pending := make(map[int64]Entry, len(listed)) + ahead := false + for _, entry := range listed { + seq, ok := parseEntry(entry.Name) + if !ok { + continue + } + pending[seq] = entry + if seq > r.seq { + ahead = true + } + } + + if _, ok := pending[r.seq]; !ok && ahead { + if r.holeSince.IsZero() { + r.holeSince = time.Now() + } else if time.Since(r.holeSince) >= r.holeTimeout { + return false, false, errors.New("segment ", r.seq, + " never arrived while later ones did, the peer lost it") + } + } else { + r.holeSince = time.Time{} + } + + for { + if entry, ok := pending[r.seq]; ok && strings.HasSuffix(entry.Name, errSuffix) { + r.discard(r.prefix + "/" + entry.Name) + return advanced, false, errors.New("the peer could not store segment ", r.seq) + } + + batch, done := r.nextBatch(pending) + if done { + return advanced, true, nil + } + if len(batch) == 0 { + return advanced, false, nil + } + + chunks, err := r.fetch(batch) + if err != nil { + if err == errNotFound { + return advanced, false, nil + } + return advanced, false, err + } + + for i, chunk := range chunks { + select { + case r.ch <- chunk: + case <-r.ctx.Done(): + return advanced, true, nil + } + r.seq++ + advanced = true + r.discard(r.prefix + "/" + batch[i].Name) + } + } +} + +func (r *walReader) nextBatch(pending map[int64]Entry) (batch []Entry, done bool) { + for i := 0; i < r.concurrency; i++ { + entry, ok := pending[r.seq+int64(i)] + if !ok { + break + } + if !strings.HasSuffix(entry.Name, segSuffix) { + if i == 0 && strings.HasSuffix(entry.Name, endSuffix) { + r.discard(r.prefix + "/" + entry.Name) + return nil, true + } + break + } + batch = append(batch, entry) + } + return batch, false +} + +func (r *walReader) fetch(batch []Entry) ([][]byte, error) { + chunks := make([][]byte, len(batch)) + failures := make([]error, len(batch)) + + var wg sync.WaitGroup + for i, entry := range batch { + if entry.Inline != nil { + chunks[i] = entry.Inline + continue + } + wg.Add(1) + go func(i int, name string) { + defer wg.Done() + chunks[i], failures[i] = r.storage.Get(r.ctx, r.prefix+"/"+name) + }(i, entry.Name) + } + wg.Wait() + + for i := range batch { + if failures[i] != nil { + return nil, failures[i] + } + } + return chunks, nil +} + +func (r *walReader) discard(name string) { + select { + case r.discards <- name: + default: + } +} + +func (r *walReader) discardLoop() { + var wg sync.WaitGroup + defer wg.Wait() + + sem := make(chan struct{}, r.concurrency) + for { + select { + case <-r.ctx.Done(): + return + case name := <-r.discards: + sem <- struct{}{} + wg.Add(1) + go func(name string) { + defer wg.Done() + defer func() { <-sem }() + r.storage.Delete(r.ctx, name) + }(name) + } + } +} + +func (r *walReader) setErr(err error) { + r.errMu.Lock() + defer r.errMu.Unlock() + if r.err == nil { + r.err = err + } +} + +func (r *walReader) Err() error { + r.errMu.Lock() + defer r.errMu.Unlock() + if r.err != nil { + return r.err + } + return io.EOF +} diff --git a/transport/internet/xdrive/xdrive.go b/transport/internet/xdrive/xdrive.go new file mode 100644 index 000000000..3c078b252 --- /dev/null +++ b/transport/internet/xdrive/xdrive.go @@ -0,0 +1,308 @@ +package xdrive + +import ( + "context" + "crypto/rand" + "encoding/hex" + "fmt" + "strconv" + "strings" + "sync" + "time" + + "github.com/xtls/xray-core/common" + "github.com/xtls/xray-core/common/errors" + "github.com/xtls/xray-core/common/net" + "github.com/xtls/xray-core/transport/internet" + "github.com/xtls/xray-core/transport/internet/stat" +) + +const ( + protocolName = "xdrive" + sessionsDir = "sessions" + streamsDir = "streams" + uplinkDir = "c2s" + downlinkDir = "s2c" +) + +func init() { + common.Must(internet.RegisterProtocolConfigCreator(protocolName, func() interface{} { + return new(Config) + })) + common.Must(internet.RegisterTransportDialer(protocolName, Dial)) + common.Must(internet.RegisterTransportListener(protocolName, Serve)) +} + +func newSessionID() (string, error) { + buf := make([]byte, 16) + if _, err := rand.Read(buf); err != nil { + return "", errors.New("failed to generate session id").Base(err) + } + return hex.EncodeToString(buf), nil +} + +func announceName(session string, at time.Time) string { + return fmt.Sprintf("%s/%d-%s", sessionsDir, at.UnixNano(), session) +} + +func parseAnnounce(entry string) (string, time.Time, bool) { + dash := strings.IndexByte(entry, '-') + if dash <= 0 || dash == len(entry)-1 { + return "", time.Time{}, false + } + nanos, err := strconv.ParseInt(entry[:dash], 10, 64) + if err != nil { + return "", time.Time{}, false + } + return entry[dash+1:], time.Unix(0, nanos), true +} + +func sessionPrefix(session string) string { + return streamsDir + "/" + session +} + +func uplinkPrefix(session string) string { + return sessionPrefix(session) + "/" + uplinkDir +} + +func downlinkPrefix(session string) string { + return sessionPrefix(session) + "/" + downlinkDir +} + +func streamConfig(streamSettings *internet.MemoryStreamConfig) (*Config, error) { + config, ok := streamSettings.ProtocolSettings.(*Config) + if !ok || config == nil { + return nil, errors.New("invalid protocol settings") + } + return config, nil +} + +func Dial(ctx context.Context, dest net.Destination, streamSettings *internet.MemoryStreamConfig) (stat.Connection, error) { + config, err := streamConfig(streamSettings) + if err != nil { + return nil, err + } + + storage, err := newStorage(streamSettings) + if err != nil { + return nil, err + } + + session, err := newSessionID() + if err != nil { + storage.Close() + return nil, err + } + + if err := storage.Put(ctx, announceName(session, time.Now()), nil); err != nil { + storage.Close() + return nil, errors.New("failed to announce session ", session).Base(err) + } + + errors.LogInfo(ctx, "opened session ", session) + + return newConn(context.Background(), storage, + uplinkPrefix(session), downlinkPrefix(session), paramsFromConfig(config), func() { + storage.Close() + }), nil +} + +type Listener struct { + ctx context.Context + cancel context.CancelFunc + storage Storage + addConn internet.ConnHandler + params + + mu sync.Mutex + active map[string]bool + handled map[string]time.Time + idleSince map[string]time.Time +} + +func Serve(ctx context.Context, address net.Address, port net.Port, streamSettings *internet.MemoryStreamConfig, addConn internet.ConnHandler) (internet.Listener, error) { + config, err := streamConfig(streamSettings) + if err != nil { + return nil, err + } + + storage, err := newStorage(streamSettings) + if err != nil { + return nil, err + } + + listenerCtx, cancel := context.WithCancel(context.Background()) + listener := &Listener{ + ctx: listenerCtx, + cancel: cancel, + storage: storage, + addConn: addConn, + params: paramsFromConfig(config), + active: make(map[string]bool), + handled: make(map[string]time.Time), + idleSince: make(map[string]time.Time), + } + + go listener.acceptLoop(ctx) + go listener.collectLoop(ctx) + + return listener, nil +} + +func (l *Listener) acceptLoop(logCtx context.Context) { + delay := l.minPollInterval + active := time.Now() + for { + accepted, err := l.acceptPending(logCtx) + if err != nil { + errors.LogWarningInner(logCtx, err, "failed to list sessions") + } + + switch { + case accepted: + active = time.Now() + delay = l.minPollInterval + case time.Since(active) < l.eagerWindow: + delay = l.minPollInterval + default: + delay *= 2 + if delay > l.maxPollInterval { + delay = l.maxPollInterval + } + } + + select { + case <-l.ctx.Done(): + return + case <-time.After(delay): + } + } +} + +func (l *Listener) acceptPending(logCtx context.Context) (bool, error) { + sessions, err := l.storage.List(l.ctx, sessionsDir) + if err != nil { + return false, err + } + + accepted := false + for _, listed := range sessions { + entry := listed.Name + full := sessionsDir + "/" + entry + + session, at, ok := parseAnnounce(entry) + if !ok { + go l.drop(full) + continue + } + if time.Since(at) > l.sessionTTL { + errors.LogInfo(logCtx, "dropping the stale announcement of session ", session) + go l.drop(full) + go l.drop(sessionPrefix(session)) + continue + } + if !l.claim(session) { + continue + } + go l.drop(full) + errors.LogInfo(logCtx, "accepted session ", session) + accepted = true + l.addConn(l.newSessionConn(session)) + } + return accepted, nil +} + +func (l *Listener) drop(name string) { + l.storage.Delete(l.ctx, name) +} + +func (l *Listener) claim(session string) bool { + l.mu.Lock() + defer l.mu.Unlock() + if l.active[session] || !l.handled[session].IsZero() { + return false + } + l.active[session] = true + l.handled[session] = time.Now() + return true +} + +func (l *Listener) newSessionConn(session string) *Conn { + return newConn(l.ctx, l.storage, + downlinkPrefix(session), uplinkPrefix(session), l.params, func() { + l.mu.Lock() + delete(l.active, session) + l.mu.Unlock() + }) +} + +func (l *Listener) collectLoop(logCtx context.Context) { + interval := l.sessionTTL / 2 + for { + select { + case <-l.ctx.Done(): + return + case <-time.After(interval): + } + if err := l.collect(); err != nil { + errors.LogWarningInner(logCtx, err, "failed to collect abandoned sessions") + } + } +} + +func (l *Listener) collect() error { + sessions, err := l.storage.List(l.ctx, streamsDir) + if err != nil { + return err + } + + now := time.Now() + var expired []string + + l.mu.Lock() + present := make(map[string]bool, len(sessions)) + for _, listed := range sessions { + session := listed.Name + present[session] = true + if l.active[session] { + delete(l.idleSince, session) + continue + } + since, seen := l.idleSince[session] + if !seen { + l.idleSince[session] = now + continue + } + if now.Sub(since) >= l.sessionTTL { + expired = append(expired, session) + delete(l.idleSince, session) + } + } + for session := range l.idleSince { + if !present[session] { + delete(l.idleSince, session) + } + } + for session, at := range l.handled { + if !l.active[session] && now.Sub(at) >= l.sessionTTL { + delete(l.handled, session) + } + } + l.mu.Unlock() + + for _, session := range expired { + if err := l.storage.Delete(l.ctx, sessionPrefix(session)); err != nil { + return err + } + } + return nil +} + +func (l *Listener) Addr() net.Addr { + return placeholderAddr +} + +func (l *Listener) Close() error { + l.cancel() + return l.storage.Close() +} diff --git a/transport/internet/xdrive/xdrive_test.go b/transport/internet/xdrive/xdrive_test.go new file mode 100644 index 000000000..b91247f76 --- /dev/null +++ b/transport/internet/xdrive/xdrive_test.go @@ -0,0 +1,583 @@ +package xdrive + +import ( + "bytes" + "context" + "crypto/rand" + "io" + "os" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/xtls/xray-core/common/net" + "github.com/xtls/xray-core/transport/internet" + "github.com/xtls/xray-core/transport/internet/stat" +) + +const testPatience = 30 * time.Second + +func settings(folder string) *internet.MemoryStreamConfig { + return &internet.MemoryStreamConfig{ + ProtocolName: protocolName, + ProtocolSettings: &Config{ + RemoteFolder: folder, + Service: "local", + FlushIntervalMs: 5, + PollIntervalMs: 5, + MaxPollIntervalMs: 20, + SessionTtlSeconds: 5, + }, + } +} + +func pair(t *testing.T) (client, server stat.Connection, cleanup func()) { + t.Helper() + return pairWith(t, settings(t.TempDir())) +} + +func pairWith(t *testing.T, streamSettings *internet.MemoryStreamConfig) (client, server stat.Connection, cleanup func()) { + t.Helper() + + accepted := make(chan stat.Connection, 1) + + listener, err := Serve(context.Background(), net.LocalHostIP, net.Port(0), streamSettings, func(conn stat.Connection) { + accepted <- conn + }) + if err != nil { + t.Fatalf("Serve: %v", err) + } + + client, err = Dial(context.Background(), net.Destination{}, streamSettings) + if err != nil { + listener.Close() + t.Fatalf("Dial: %v", err) + } + + select { + case server = <-accepted: + case <-time.After(testPatience): + client.Close() + listener.Close() + t.Fatal("listener did not accept the session") + } + + return client, server, func() { + client.Close() + server.Close() + listener.Close() + } +} + +func expectRead(t *testing.T, conn stat.Connection, want string) { + t.Helper() + + if err := conn.SetReadDeadline(time.Now().Add(testPatience)); err != nil { + t.Fatalf("SetReadDeadline: %v", err) + } + buf := make([]byte, len(want)) + if _, err := io.ReadFull(conn, buf); err != nil { + t.Fatalf("ReadFull: %v", err) + } + if string(buf) != want { + t.Fatalf("read %q, want %q", buf, want) + } +} + +func TestRoundTrip(t *testing.T) { + client, server, cleanup := pair(t) + defer cleanup() + + if _, err := client.Write([]byte("ping")); err != nil { + t.Fatalf("client write: %v", err) + } + expectRead(t, server, "ping") + + if _, err := server.Write([]byte("pong")); err != nil { + t.Fatalf("server write: %v", err) + } + expectRead(t, client, "pong") +} + +func TestInterleaved(t *testing.T) { + client, server, cleanup := pair(t) + defer cleanup() + + for i := 0; i < 20; i++ { + if _, err := client.Write([]byte("up")); err != nil { + t.Fatalf("client write %d: %v", i, err) + } + expectRead(t, server, "up") + + if _, err := server.Write([]byte("down")); err != nil { + t.Fatalf("server write %d: %v", i, err) + } + expectRead(t, client, "down") + } +} + +func TestMultiSegmentTransfer(t *testing.T) { + client, server, cleanup := pair(t) + defer cleanup() + + payload := make([]byte, 3*defaultSegmentBytes+1234) + if _, err := rand.Read(payload); err != nil { + t.Fatalf("rand: %v", err) + } + + go func() { + client.Write(payload) + }() + + if err := server.SetReadDeadline(time.Now().Add(30 * time.Second)); err != nil { + t.Fatalf("SetReadDeadline: %v", err) + } + got := make([]byte, len(payload)) + if _, err := io.ReadFull(server, got); err != nil { + t.Fatalf("ReadFull: %v", err) + } + if !bytes.Equal(got, payload) { + t.Fatal("payload mismatch") + } +} + +func TestCloseEOF(t *testing.T) { + client, server, cleanup := pair(t) + defer cleanup() + + if _, err := client.Write([]byte("bye")); err != nil { + t.Fatalf("client write: %v", err) + } + if err := client.Close(); err != nil { + t.Fatalf("client close: %v", err) + } + + if err := server.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatalf("SetReadDeadline: %v", err) + } + got, err := io.ReadAll(server) + if err != nil { + t.Fatalf("ReadAll: %v", err) + } + if string(got) != "bye" { + t.Fatalf("read %q, want %q", got, "bye") + } +} + +func TestReadDeadline(t *testing.T) { + client, _, cleanup := pair(t) + defer cleanup() + + if err := client.SetReadDeadline(time.Now().Add(100 * time.Millisecond)); err != nil { + t.Fatalf("SetReadDeadline: %v", err) + } + buf := make([]byte, 4) + if _, err := client.Read(buf); !os.IsTimeout(err) { + t.Fatalf("Read returned %v, want a timeout", err) + } +} + +func TestLocalNameEscape(t *testing.T) { + root := t.TempDir() + storage, err := newLocalStorage(root) + if err != nil { + t.Fatalf("newLocalStorage: %v", err) + } + + if err := storage.Put(context.Background(), "../escaped", []byte("x")); err != nil { + t.Fatalf("Put: %v", err) + } + if _, err := os.Stat(root + "/escaped"); err != nil { + t.Fatalf("name was not clamped inside the root: %v", err) + } +} + +func TestLocalMissingObject(t *testing.T) { + storage, err := newLocalStorage(t.TempDir()) + if err != nil { + t.Fatalf("newLocalStorage: %v", err) + } + + if _, err := storage.Get(context.Background(), "nothing/here"); err != errNotFound { + t.Fatalf("Get returned %v, want errNotFound", err) + } + names, err := storage.List(context.Background(), "nothing") + if err != nil { + t.Fatalf("List: %v", err) + } + if len(names) != 0 { + t.Fatalf("List returned %v, want none", names) + } +} + +func TestResumeAfterIdle(t *testing.T) { + client, server, cleanup := pair(t) + defer cleanup() + + if _, err := client.Write([]byte("first")); err != nil { + t.Fatalf("client write: %v", err) + } + expectRead(t, server, "first") + + time.Sleep(200 * time.Millisecond) + + if _, err := client.Write([]byte("second")); err != nil { + t.Fatalf("client write: %v", err) + } + expectRead(t, server, "second") +} + +func TestParseEntry(t *testing.T) { + cases := []struct { + name string + seq int64 + ok bool + }{ + {"000000000.seg", 0, true}, + {"000000042.seg", 42, true}, + {"000000007.end", 7, true}, + {"000000001.tmp", 0, false}, + {"notanumber.seg", 0, false}, + {"000000001", 0, false}, + } + for _, c := range cases { + seq, ok := parseEntry(c.name) + if ok != c.ok || (ok && seq != c.seq) { + t.Fatalf("parseEntry(%q) = %d, %v; want %d, %v", c.name, seq, ok, c.seq, c.ok) + } + } +} + +func TestParamDefaults(t *testing.T) { + p := paramsFromConfig(&Config{}) + if p.segmentBytes != defaultSegmentBytes || p.flushInterval != defaultFlushInterval { + t.Fatalf("defaults not applied: %+v", p) + } + + p = paramsFromConfig(&Config{SegmentBytes: 1 << 30, PollIntervalMs: 400, MaxPollIntervalMs: 100}) + if p.segmentBytes != maxSegmentBytes { + t.Fatalf("segmentBytes is %d, want %d", p.segmentBytes, maxSegmentBytes) + } + if p.maxPollInterval < p.minPollInterval { + t.Fatalf("maxPollInterval %v below minPollInterval %v", p.maxPollInterval, p.minPollInterval) + } +} + +func waitFor(t *testing.T, what string, done func() bool) { + t.Helper() + + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if done() { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("timed out waiting for %s", what) +} + +func newTestListener(t *testing.T, folder string) *Listener { + t.Helper() + + storage, err := newLocalStorage(folder) + if err != nil { + t.Fatalf("newLocalStorage: %v", err) + } + return &Listener{ + ctx: context.Background(), + storage: storage, + params: paramsFromConfig(&Config{SessionTtlSeconds: 1}), + active: make(map[string]bool), + handled: make(map[string]time.Time), + idleSince: make(map[string]time.Time), + } +} + +func TestCollectAbandoned(t *testing.T) { + folder := t.TempDir() + listener := newTestListener(t, folder) + + if err := listener.storage.Put(context.Background(), uplinkPrefix("dead")+"/000000000.seg", []byte("x")); err != nil { + t.Fatalf("Put: %v", err) + } + + if err := listener.collect(); err != nil { + t.Fatalf("collect: %v", err) + } + names, _ := listener.storage.List(context.Background(), streamsDir) + if len(names) != 1 { + t.Fatalf("first pass removed the session, got %v", names) + } + + listener.idleSince["dead"] = time.Now().Add(-2 * time.Second) + if err := listener.collect(); err != nil { + t.Fatalf("collect: %v", err) + } + names, _ = listener.storage.List(context.Background(), streamsDir) + if len(names) != 0 { + t.Fatalf("abandoned session still there, got %v", names) + } +} + +func TestCollectKeepsActive(t *testing.T) { + folder := t.TempDir() + listener := newTestListener(t, folder) + listener.active["live"] = true + + if err := listener.storage.Put(context.Background(), uplinkPrefix("live")+"/000000000.seg", []byte("x")); err != nil { + t.Fatalf("Put: %v", err) + } + + listener.idleSince["live"] = time.Now().Add(-2 * time.Second) + if err := listener.collect(); err != nil { + t.Fatalf("collect: %v", err) + } + names, _ := listener.storage.List(context.Background(), streamsDir) + if len(names) != 1 { + t.Fatalf("collected an active session, got %v", names) + } +} + +func TestParseAnnounce(t *testing.T) { + session, at, ok := parseAnnounce("1757000000123456789-abc123") + if !ok || session != "abc123" || at.UnixNano() != 1757000000123456789 { + t.Fatalf("parseAnnounce returned %q, %v, %v", session, at.UnixNano(), ok) + } + for _, bad := range []string{"abc123", "-abc123", "1757000000-", "notanumber-abc"} { + if _, _, ok := parseAnnounce(bad); ok { + t.Fatalf("parseAnnounce accepted %q", bad) + } + } +} + +func TestStaleAnnounce(t *testing.T) { + folder := t.TempDir() + listener := newTestListener(t, folder) + + ctx := context.Background() + stale := announceName("ghost", time.Now().Add(-time.Hour)) + if err := listener.storage.Put(ctx, stale, nil); err != nil { + t.Fatalf("Put: %v", err) + } + if err := listener.storage.Put(ctx, uplinkPrefix("ghost")+"/000000000.seg", []byte("x")); err != nil { + t.Fatalf("Put: %v", err) + } + + accepted, err := listener.acceptPending(ctx) + if err != nil { + t.Fatalf("acceptPending: %v", err) + } + if accepted { + t.Fatal("accepted a stale announcement") + } + + waitFor(t, "the stale announcement to be removed", func() bool { + names, _ := listener.storage.List(ctx, sessionsDir) + return len(names) == 0 + }) + waitFor(t, "the stale session data to be removed", func() bool { + names, _ := listener.storage.List(ctx, streamsDir) + return len(names) == 0 + }) +} + +func TestFreshAnnounce(t *testing.T) { + folder := t.TempDir() + listener := newTestListener(t, folder) + listener.addConn = func(conn stat.Connection) { conn.Close() } + + ctx := context.Background() + if err := listener.storage.Put(ctx, announceName("fresh", time.Now()), nil); err != nil { + t.Fatalf("Put: %v", err) + } + + accepted, err := listener.acceptPending(ctx) + if err != nil { + t.Fatalf("acceptPending: %v", err) + } + if !accepted { + t.Fatal("did not accept a fresh announcement") + } +} + +func TestAnnouncePrecision(t *testing.T) { + at := time.Unix(1757000000, int64(900*time.Millisecond)) + entry := strings.TrimPrefix(announceName("abc123", at), sessionsDir+"/") + + session, parsed, ok := parseAnnounce(entry) + if !ok || session != "abc123" { + t.Fatalf("parseAnnounce(%q) returned %q, %v", entry, session, ok) + } + if !parsed.Equal(at) { + t.Fatalf("timestamp came back as %v, want %v", parsed, at) + } +} + +func TestRecentAnnounceTTL(t *testing.T) { + folder := t.TempDir() + listener := newTestListener(t, folder) + listener.addConn = func(conn stat.Connection) { conn.Close() } + + ctx := context.Background() + recent := time.Now().Add(-900 * time.Millisecond) + if err := listener.storage.Put(ctx, announceName("recent", recent), nil); err != nil { + t.Fatalf("Put: %v", err) + } + + accepted, err := listener.acceptPending(ctx) + if err != nil { + t.Fatalf("acceptPending: %v", err) + } + if !accepted { + t.Fatal("dropped an announcement younger than the TTL") + } +} + +func TestMissingSegment(t *testing.T) { + storage, err := newLocalStorage(t.TempDir()) + if err != nil { + t.Fatalf("newLocalStorage: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + p := paramsFromConfig(&Config{PollIntervalMs: 5, MaxPollIntervalMs: 20, HoleTimeoutMs: 200}) + if err := storage.Put(ctx, objectName("hole", 1, segSuffix), []byte("second")); err != nil { + t.Fatalf("Put: %v", err) + } + + reader := newWALReader(ctx, storage, "hole", p) + select { + case _, ok := <-reader.ch: + if ok { + t.Fatal("delivered data past a missing segment") + } + case <-time.After(5 * time.Second): + t.Fatal("reader did not give up on a missing segment") + } + + err = reader.Err() + if err == nil || err == io.EOF { + t.Fatalf("Err returned %v, want a failure", err) + } +} + +func TestIdleStreamWaits(t *testing.T) { + storage, err := newLocalStorage(t.TempDir()) + if err != nil { + t.Fatalf("newLocalStorage: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + p := paramsFromConfig(&Config{PollIntervalMs: 5, MaxPollIntervalMs: 20, HoleTimeoutMs: 100}) + reader := newWALReader(ctx, storage, "idle", p) + + time.Sleep(400 * time.Millisecond) + if err := storage.Put(ctx, objectName("idle", 0, segSuffix), []byte("late")); err != nil { + t.Fatalf("Put: %v", err) + } + + select { + case data, ok := <-reader.ch: + if !ok { + t.Fatalf("the reader gave up on an idle stream: %v", reader.Err()) + } + if string(data) != "late" { + t.Fatalf("read %q, want %q", data, "late") + } + case <-time.After(5 * time.Second): + t.Fatal("reader missed a late segment") + } +} + +func TestFailureMarker(t *testing.T) { + storage, err := newLocalStorage(t.TempDir()) + if err != nil { + t.Fatalf("newLocalStorage: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + p := paramsFromConfig(&Config{PollIntervalMs: 5, MaxPollIntervalMs: 20}) + if err := storage.Put(ctx, objectName("broken", 0, segSuffix), []byte("first")); err != nil { + t.Fatalf("Put: %v", err) + } + if err := storage.Put(ctx, objectName("broken", 1, errSuffix), nil); err != nil { + t.Fatalf("Put: %v", err) + } + + reader := newWALReader(ctx, storage, "broken", p) + + select { + case data, ok := <-reader.ch: + if !ok || string(data) != "first" { + t.Fatalf("want the segment before the marker, got %q %v", data, ok) + } + case <-time.After(5 * time.Second): + t.Fatal("reader did not deliver the first segment") + } + + select { + case _, ok := <-reader.ch: + if ok { + t.Fatal("delivered data past the failure marker") + } + case <-time.After(5 * time.Second): + t.Fatal("reader did not stop on the failure marker") + } + + if err := reader.Err(); err == nil || err == io.EOF { + t.Fatalf("Err returned %v, want a failure", err) + } +} + +type inlineOnlyStorage struct { + Storage + gets int64 +} + +func (s *inlineOnlyStorage) List(ctx context.Context, prefix string) ([]Entry, error) { + return []Entry{{Name: "000000000" + segSuffix, Inline: []byte("carried by the listing")}}, nil +} + +func (s *inlineOnlyStorage) Get(ctx context.Context, name string) ([]byte, error) { + atomic.AddInt64(&s.gets, 1) + return nil, errNotFound +} + +func (s *inlineOnlyStorage) Delete(ctx context.Context, name string) error { + return nil +} + +func TestInlinePayload(t *testing.T) { + base, err := newLocalStorage(t.TempDir()) + if err != nil { + t.Fatalf("newLocalStorage: %v", err) + } + storage := &inlineOnlyStorage{Storage: base} + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + reader := newWALReader(ctx, storage, "inline", paramsFromConfig(&Config{PollIntervalMs: 5})) + select { + case data, ok := <-reader.ch: + if !ok { + t.Fatalf("reader stopped: %v", reader.Err()) + } + if string(data) != "carried by the listing" { + t.Fatalf("read %q, want the inline payload", data) + } + case <-time.After(5 * time.Second): + t.Fatal("reader did not deliver the inline payload") + } + + if got := atomic.LoadInt64(&storage.gets); got != 0 { + t.Fatalf("called Get %d times for an inline payload", got) + } +}