Compare commits

...
2 Commits
22 changed files with 4960 additions and 0 deletions
+13
View File
@@ -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 {
+48
View File
@@ -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"
)
@@ -794,3 +795,50 @@ func readFileOrString(f string, s []string) ([]byte, error) {
}
return nil, errors.New("both file and bytes are empty.")
}
type XDriveConfig struct {
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"`
Template json.RawMessage `json:"template"`
}
// 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")
}
case "template":
if len(c.Template) == 0 {
return nil, errors.New(`service "template" needs a "template" object`)
}
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,
Template: string(c.Template),
}
return config, nil
}
+73
View File
@@ -291,3 +291,76 @@ 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")
}
}
func TestXDriveTemplateStreamConfig(t *testing.T) {
config := new(StreamConfig)
if err := json.Unmarshal([]byte(`{
"method": "xdrive",
"xdriveSettings": {
"remoteFolder": "folder",
"service": "template",
"secrets": ["user", "pass"],
"template": {
"flatten": true,
"auth": {"type": "basic", "username": "{secret0}", "password": "{secret1}"},
"put": {"method": "PUT", "url": "https://dav.example/{folder}/{name}"},
"get": {"method": "GET", "url": "https://dav.example/{folder}/{name}"},
"delete": {"method": "DELETE", "url": "https://dav.example/{folder}/{name}"},
"list": {"method": "PROPFIND", "url": "https://dav.example/{folder}/", "namesRegex": "<d:href>/folder/([^<]+)</d:href>"}
}
}
}`), 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 xdrive", built.ProtocolName)
}
}
func TestXDriveTemplateNeedsTemplate(t *testing.T) {
config := new(XDriveConfig)
if err := json.Unmarshal([]byte(`{"remoteFolder": "f", "service": "template"}`), config); err != nil {
t.Fatalf("Unmarshal: %v", err)
}
if _, err := config.Build(); err == nil {
t.Fatal("Build accepted a template service without a template")
}
}
+1
View File
@@ -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"
+156
View File
@@ -0,0 +1,156 @@
package xdrive
import (
"context"
gotls "crypto/tls"
"net/http"
"time"
"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/reality"
"github.com/xtls/xray-core/transport/internet/tls"
"golang.org/x/net/http2"
)
type serviceTransport struct {
plain http.RoundTripper
secure http.RoundTripper
}
func (t *serviceTransport) RoundTrip(r *http.Request) (*http.Response, error) {
if r.URL.Scheme == "https" {
return t.secure.RoundTrip(r)
}
return t.plain.RoundTrip(r)
}
func newServiceClient(streamSettings *internet.MemoryStreamConfig, timeout time.Duration, maxConns int) *http.Client {
var (
tlsConfig *tls.Config
realityConfig *reality.Config
sockopt *internet.SocketConfig
fronting *net.Destination
)
if streamSettings != nil {
tlsConfig = tls.ConfigFromStreamSettings(streamSettings)
realityConfig = reality.ConfigFromStreamSettings(streamSettings)
sockopt = streamSettings.SocketSettings
fronting = streamSettings.Destination
}
overHTTP2 := allowsHTTP2(tlsConfig, realityConfig)
dial := func(ctx context.Context, addr string) (net.Conn, net.Destination, error) {
host, err := net.ParseDestination("tcp:" + addr)
if err != nil {
return nil, host, errors.New("bad address: ", addr).Base(err)
}
target := host
if fronting != nil {
target.Address = fronting.Address
if fronting.Port != 0 {
target.Port = fronting.Port
}
}
conn, err := internet.DialSystem(ctx, target, sockopt)
if err != nil {
return nil, host, err
}
if streamSettings != nil && streamSettings.TcpmaskManager != nil {
masked, err := streamSettings.TcpmaskManager.WrapConnClient(conn)
if err != nil {
conn.Close()
return nil, host, errors.New("mask err").Base(err)
}
conn = masked
}
return conn, host, nil
}
dialPlain := func(ctx context.Context, network, addr string) (net.Conn, error) {
conn, _, err := dial(ctx, addr)
return conn, err
}
dialTLS := func(ctx context.Context, addr string) (net.Conn, error) {
conn, host, err := dial(ctx, addr)
if err != nil {
return nil, err
}
if realityConfig != nil {
return reality.UClient(conn, realityConfig, ctx, host)
}
gotlsConfig := &gotls.Config{ServerName: host.Address.String()}
if tlsConfig != nil {
gotlsConfig = tlsConfig.GetTLSConfig(tls.WithDestination(host))
}
if len(gotlsConfig.NextProtos) != 1 {
if overHTTP2 {
gotlsConfig.NextProtos = []string{"h2"}
} else {
gotlsConfig.NextProtos = []string{"http/1.1"}
}
}
if tlsConfig != nil {
if fingerprint := tls.GetFingerprint(tlsConfig.Fingerprint); fingerprint != nil {
uconn := tls.UClient(conn, gotlsConfig, fingerprint)
if err := uconn.(*tls.UConn).HandshakeContext(ctx); err != nil {
conn.Close()
return nil, err
}
return uconn, nil
}
}
return tls.Client(conn, gotlsConfig), nil
}
var secure http.RoundTripper
if overHTTP2 {
secure = &http2.Transport{
DialTLSContext: func(ctx context.Context, network, addr string, cfg *gotls.Config) (net.Conn, error) {
return dialTLS(ctx, addr)
},
IdleConnTimeout: net.ConnIdleTimeout,
}
} else {
secure = &http.Transport{
DialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
return dialTLS(ctx, addr)
},
IdleConnTimeout: net.ConnIdleTimeout,
MaxIdleConns: maxConns,
MaxIdleConnsPerHost: maxConns,
MaxConnsPerHost: maxConns,
}
}
return &http.Client{
Transport: &serviceTransport{
plain: &http.Transport{
DialContext: dialPlain,
IdleConnTimeout: net.ConnIdleTimeout,
MaxIdleConns: maxConns,
MaxIdleConnsPerHost: maxConns,
MaxConnsPerHost: maxConns,
},
secure: secure,
},
Timeout: timeout,
}
}
func allowsHTTP2(tlsConfig *tls.Config, realityConfig *reality.Config) bool {
if realityConfig != nil {
return true
}
if tlsConfig == nil {
return false
}
return len(tlsConfig.NextProtocol) == 1 && tlsConfig.NextProtocol[0] == "h2"
}
+101
View File
@@ -0,0 +1,101 @@
package xdrive
import (
"context"
gotls "crypto/tls"
"net"
"net/http"
"sync"
"testing"
"time"
xnet "github.com/xtls/xray-core/common/net"
"github.com/xtls/xray-core/transport/internet"
"github.com/xtls/xray-core/transport/internet/tls"
)
func recordingTLSListener(t *testing.T, sni *string, mu *sync.Mutex) net.Listener {
t.Helper()
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
go func() {
for {
conn, err := ln.Accept()
if err != nil {
return
}
cfg := &gotls.Config{
GetConfigForClient: func(hello *gotls.ClientHelloInfo) (*gotls.Config, error) {
mu.Lock()
*sni = hello.ServerName
mu.Unlock()
return nil, nil
},
}
tconn := gotls.Server(conn, cfg)
tconn.HandshakeContext(context.Background())
tconn.Close()
}
}()
t.Cleanup(func() { ln.Close() })
return ln
}
func sniForSettings(t *testing.T, serverName string) string {
t.Helper()
var (
sni string
mu sync.Mutex
)
ln := recordingTLSListener(t, &sni, &mu)
addr := ln.Addr().(*net.TCPAddr)
settings := &internet.MemoryStreamConfig{
ProtocolName: protocolName,
Destination: &xnet.Destination{
Address: xnet.ParseAddress(addr.IP.String()),
Port: xnet.Port(addr.Port),
Network: xnet.Network_TCP,
},
SecuritySettings: &tls.Config{ServerName: serverName},
}
prev := driveFilesURL
driveFilesURL = "https://www.googleapis.com/drive/v3/files"
defer func() { driveFilesURL = prev }()
client := newServiceClient(settings, 5*time.Second, 8)
req, err := http.NewRequest(http.MethodGet, driveFilesURL, nil)
if err != nil {
t.Fatalf("request: %v", err)
}
client.Do(req)
deadline := time.Now().Add(3 * time.Second)
for time.Now().Before(deadline) {
mu.Lock()
got := sni
mu.Unlock()
if got != "" {
return got
}
time.Sleep(5 * time.Millisecond)
}
return ""
}
func TestServiceSNIDefaultsToHost(t *testing.T) {
if got := sniForSettings(t, ""); got != "www.googleapis.com" {
t.Fatalf("SNI defaulted to %q, want the host www.googleapis.com, not address", got)
}
}
func TestServiceSNIOverride(t *testing.T) {
if got := sniForSettings(t, "www.google.com"); got != "www.google.com" {
t.Fatalf("explicit serverName gave SNI %q, want www.google.com", got)
}
}
+222
View File
@@ -0,0 +1,222 @@
// 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"`
Template string `protobuf:"bytes,12,opt,name=template,proto3" json:"template,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
}
func (x *Config) GetTemplate() string {
if x != nil {
return x.Template
}
return ""
}
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\"\xcb\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\rholeTimeoutMs\x12\x1a\n" +
"\btemplate\x18\f \x01(\tR\btemplateB5Z3github.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
}
+19
View File
@@ -0,0 +1,19 @@
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;
string template = 12;
}
+139
View File
@@ -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
}
+575
View File
@@ -0,0 +1,575 @@
package xdrive
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"sync"
"time"
"github.com/xtls/xray-core/common/dice"
"github.com/xtls/xray-core/common/errors"
"github.com/xtls/xray-core/transport/internet"
)
const (
flatSeparator = "~"
driveBoundary = "xdrive-boundary"
drivePageSize = 1000
driveMaxAttempts = 8
driveMaxInflight = 32
driveInlineLimit = 12000
driveTimeout = 60 * time.Second
)
var (
driveMaxBackoff = 8 * time.Second
driveTokenURL = "https://oauth2.googleapis.com/token"
driveFilesURL = "https://www.googleapis.com/drive/v3/files"
driveUploadURL = "https://www.googleapis.com/upload/drive/v3/files?uploadType=multipart&fields=id,name"
driveInitialBackoff = 200 * time.Millisecond
)
type driveStorage struct {
folder string
clientID string
clientSecret string
refreshToken string
client *http.Client
tokenURL string
filesURL string
uploadURL string
backoff time.Duration
tokenMu sync.Mutex
token string
tokenExpiry time.Time
inflight chan struct{}
idMu sync.Mutex
ids map[string]string
}
func newDriveStorage(streamSettings *internet.MemoryStreamConfig, config *Config) (*driveStorage, error) {
if config.RemoteFolder == "" {
return nil, errors.New(`empty "remoteFolder", it must be a Google Drive folder id`)
}
if len(config.Secrets) != 3 {
return nil, errors.New("Google Drive needs 3 secrets in order of ClientID, ClientSecret, RefreshToken")
}
for i, secret := range config.Secrets {
if secret == "" {
return nil, errors.New("Google Drive secret ", i, " is empty")
}
}
return &driveStorage{
folder: config.RemoteFolder,
clientID: config.Secrets[0],
clientSecret: config.Secrets[1],
refreshToken: config.Secrets[2],
client: newServiceClient(streamSettings, driveTimeout, driveMaxInflight),
tokenURL: driveTokenURL,
filesURL: driveFilesURL,
uploadURL: driveUploadURL,
backoff: driveInitialBackoff,
inflight: make(chan struct{}, driveMaxInflight),
ids: make(map[string]string),
}, nil
}
func flatten(name string) string {
return strings.ReplaceAll(name, "/", flatSeparator)
}
func quoteDriveValue(value string) string {
return strings.NewReplacer(`\`, `\\`, `'`, `\'`).Replace(value)
}
func (s *driveStorage) accessToken(ctx context.Context) (string, error) {
s.tokenMu.Lock()
defer s.tokenMu.Unlock()
if s.token != "" && time.Now().Before(s.tokenExpiry) {
return s.token, nil
}
form := url.Values{
"client_id": {s.clientID},
"client_secret": {s.clientSecret},
"refresh_token": {s.refreshToken},
"grant_type": {"refresh_token"},
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.tokenURL, strings.NewReader(form.Encode()))
if err != nil {
return "", errors.New("failed to build the token request").Base(err)
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
resp, err := s.client.Do(req)
if err != nil {
return "", errors.New("failed to refresh the access token").Base(err)
}
defer resp.Body.Close()
body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
if err != nil {
return "", errors.New("failed to read the token response").Base(err)
}
if resp.StatusCode != http.StatusOK {
return "", errors.New("the token endpoint answered ", resp.StatusCode, ": ", string(body))
}
var parsed struct {
AccessToken string `json:"access_token"`
ExpiresIn int64 `json:"expires_in"`
}
if err := json.Unmarshal(body, &parsed); err != nil {
return "", errors.New("failed to parse the token response").Base(err)
}
if parsed.AccessToken == "" {
return "", errors.New("the token endpoint returned no access token")
}
lifetime := parsed.ExpiresIn
if lifetime > 60 {
lifetime -= 60
}
s.token = parsed.AccessToken
s.tokenExpiry = time.Now().Add(time.Duration(lifetime) * time.Second)
return s.token, nil
}
func jitter(backoff time.Duration) time.Duration {
half := backoff / 2
if half <= 0 {
return backoff
}
return half + time.Duration(dice.Roll(int(half)))
}
func rateLimited(payload []byte) bool {
var parsed struct {
Error struct {
Status string `json:"status"`
Errors []struct {
Reason string `json:"reason"`
} `json:"errors"`
} `json:"error"`
}
if json.Unmarshal(payload, &parsed) != nil {
return false
}
for _, item := range parsed.Error.Errors {
switch item.Reason {
case "rateLimitExceeded", "userRateLimitExceeded", "sharingRateLimitExceeded":
return true
}
}
return parsed.Error.Status == "RESOURCE_EXHAUSTED"
}
func retryableStatus(status int) bool {
switch status {
case http.StatusTooManyRequests, http.StatusInternalServerError,
http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout:
return true
}
return false
}
func (s *driveStorage) do(ctx context.Context, method, target, contentType string, body []byte) (int, []byte, error) {
backoff := s.backoff
var lastErr error
for attempt := 0; attempt < driveMaxAttempts; attempt++ {
if attempt > 0 {
select {
case <-ctx.Done():
return 0, nil, ctx.Err()
case <-time.After(jitter(backoff)):
}
backoff *= 2
if backoff > driveMaxBackoff {
backoff = driveMaxBackoff
}
}
token, err := s.accessToken(ctx)
if err != nil {
lastErr = err
continue
}
select {
case s.inflight <- struct{}{}:
case <-ctx.Done():
return 0, nil, ctx.Err()
}
var reader io.Reader
if body != nil {
reader = bytes.NewReader(body)
}
req, err := http.NewRequestWithContext(ctx, method, target, reader)
if err != nil {
return 0, nil, errors.New("failed to build a Drive request").Base(err)
}
req.Header.Set("Authorization", "Bearer "+token)
if contentType != "" {
req.Header.Set("Content-Type", contentType)
}
resp, err := s.client.Do(req)
<-s.inflight
if err != nil {
if ctx.Err() != nil {
return 0, nil, ctx.Err()
}
lastErr = errors.New("Drive request failed").Base(err)
errors.LogWarningInner(ctx, err, "retrying a failed Drive request, attempt ",
attempt+1, " of ", driveMaxAttempts)
continue
}
payload, err := io.ReadAll(resp.Body)
resp.Body.Close()
if err != nil {
if ctx.Err() != nil {
return 0, nil, ctx.Err()
}
lastErr = errors.New("failed to read the Drive response").Base(err)
continue
}
if resp.StatusCode == http.StatusUnauthorized {
s.invalidateToken()
lastErr = errors.New("Drive rejected the access token")
continue
}
if resp.StatusCode == http.StatusForbidden && rateLimited(payload) {
lastErr = errors.New("Drive is rate limiting: ", string(payload))
errors.LogWarning(ctx, "rate limited by Drive, attempt ",
attempt+1, " of ", driveMaxAttempts)
continue
}
if retryableStatus(resp.StatusCode) {
lastErr = errors.New("Drive answered ", resp.StatusCode, ": ", string(payload))
errors.LogWarning(ctx, "retrying after Drive answered ", resp.StatusCode,
", attempt ", attempt+1, " of ", driveMaxAttempts)
continue
}
return resp.StatusCode, payload, nil
}
return 0, nil, lastErr
}
func (s *driveStorage) invalidateToken() {
s.tokenMu.Lock()
s.token = ""
s.tokenMu.Unlock()
}
func (s *driveStorage) rememberID(name, id string) {
s.idMu.Lock()
s.ids[name] = id
s.idMu.Unlock()
}
func (s *driveStorage) forgetID(name string) {
s.idMu.Lock()
delete(s.ids, name)
s.idMu.Unlock()
}
func (s *driveStorage) cachedID(name string) (string, bool) {
s.idMu.Lock()
defer s.idMu.Unlock()
id, ok := s.ids[name]
return id, ok
}
type driveFile struct {
id string
description string
}
func (s *driveStorage) query(ctx context.Context, condition string) (map[string]driveFile, error) {
found := make(map[string]driveFile)
pageToken := ""
for {
params := url.Values{
"q": {"'" + quoteDriveValue(s.folder) + "' in parents and trashed = false and " + condition},
"fields": {"nextPageToken,files(id,name,description)"},
"pageSize": {fmt.Sprint(drivePageSize)},
"supportsAllDrives": {"true"},
"includeItemsFromAllDrives": {"true"},
}
if pageToken != "" {
params.Set("pageToken", pageToken)
}
status, payload, err := s.do(ctx, http.MethodGet, s.filesURL+"?"+params.Encode(), "", nil)
if err != nil {
return nil, err
}
if status != http.StatusOK {
return nil, errors.New("Drive listing answered ", status, ": ", string(payload))
}
var parsed struct {
NextPageToken string `json:"nextPageToken"`
Files []struct {
ID string `json:"id"`
Name string `json:"name"`
Description string `json:"description"`
} `json:"files"`
}
if err := json.Unmarshal(payload, &parsed); err != nil {
return nil, errors.New("failed to parse the Drive listing").Base(err)
}
for _, file := range parsed.Files {
found[file.Name] = driveFile{id: file.ID, description: file.Description}
s.rememberID(file.Name, file.ID)
}
pageToken = parsed.NextPageToken
if pageToken == "" {
return found, nil
}
}
}
func (s *driveStorage) resolveID(ctx context.Context, flat string) (string, error) {
if id, ok := s.cachedID(flat); ok {
return id, nil
}
found, err := s.query(ctx, "name = '"+quoteDriveValue(flat)+"'")
if err != nil {
return "", err
}
if file, ok := found[flat]; ok {
return file.id, nil
}
return "", errNotFound
}
func (s *driveStorage) Put(ctx context.Context, name string, data []byte) error {
if len(data) <= driveInlineLimit {
return s.putInline(ctx, name, data)
}
return s.putMedia(ctx, name, data)
}
func (s *driveStorage) putInline(ctx context.Context, name string, data []byte) error {
flat := flatten(name)
body, err := json.Marshal(map[string]interface{}{
"name": flat,
"parents": []string{s.folder},
"description": base64.StdEncoding.EncodeToString(data),
})
if err != nil {
return errors.New("failed to build the inline metadata").Base(err)
}
status, payload, err := s.do(ctx, http.MethodPost, s.filesURL+"?fields=id",
"application/json; charset=UTF-8", body)
if err != nil {
return err
}
if status != http.StatusOK {
return errors.New("Drive rejected the inline upload of ", name,
" with ", status, ": ", string(payload))
}
var parsed struct {
ID string `json:"id"`
}
if err := json.Unmarshal(payload, &parsed); err != nil {
return errors.New("failed to parse the Drive upload response").Base(err)
}
if parsed.ID != "" {
s.rememberID(flat, parsed.ID)
}
return nil
}
func (s *driveStorage) putMedia(ctx context.Context, name string, data []byte) error {
flat := flatten(name)
metadata, err := json.Marshal(map[string]interface{}{
"name": flat,
"parents": []string{s.folder},
})
if err != nil {
return errors.New("failed to build the upload metadata").Base(err)
}
var body bytes.Buffer
fmt.Fprintf(&body, "--%s\r\nContent-Type: application/json; charset=UTF-8\r\n\r\n", driveBoundary)
body.Write(metadata)
fmt.Fprintf(&body, "\r\n--%s\r\nContent-Type: application/octet-stream\r\n\r\n", driveBoundary)
body.Write(data)
fmt.Fprintf(&body, "\r\n--%s--\r\n", driveBoundary)
status, payload, err := s.do(ctx, http.MethodPost, s.uploadURL,
"multipart/related; boundary="+driveBoundary, body.Bytes())
if err != nil {
return err
}
if status != http.StatusOK {
return errors.New("Drive upload of ", name, " answered ", status, ": ", string(payload))
}
var parsed struct {
ID string `json:"id"`
}
if err := json.Unmarshal(payload, &parsed); err != nil {
return errors.New("failed to parse the Drive upload response").Base(err)
}
if parsed.ID != "" {
s.rememberID(flat, parsed.ID)
}
return nil
}
func (s *driveStorage) Get(ctx context.Context, name string) ([]byte, error) {
flat := flatten(name)
id, err := s.resolveID(ctx, flat)
if err != nil {
return nil, err
}
status, payload, err := s.do(ctx, http.MethodGet,
s.filesURL+"/"+url.PathEscape(id)+"?alt=media&supportsAllDrives=true", "", nil)
if err != nil {
return nil, err
}
switch status {
case http.StatusOK:
if len(payload) > 0 {
return payload, nil
}
return s.getInline(ctx, id)
case http.StatusNotFound:
s.forgetID(flat)
return nil, errNotFound
default:
return nil, errors.New("Drive download of ", name, " answered ", status, ": ", string(payload))
}
}
func (s *driveStorage) getInline(ctx context.Context, id string) ([]byte, error) {
status, payload, err := s.do(ctx, http.MethodGet,
s.filesURL+"/"+url.PathEscape(id)+"?fields=description&supportsAllDrives=true", "", nil)
if err != nil {
return nil, err
}
if status != http.StatusOK {
return nil, errors.New("Drive answered ", status, " for inline data: ", string(payload))
}
var parsed struct {
Description string `json:"description"`
}
if err := json.Unmarshal(payload, &parsed); err != nil {
return nil, errors.New("failed to parse the inline data").Base(err)
}
if parsed.Description == "" {
return nil, nil
}
data, err := base64.StdEncoding.DecodeString(parsed.Description)
if err != nil {
return nil, errors.New("the inline data is not valid base64").Base(err)
}
return data, nil
}
func (s *driveStorage) deleteID(ctx context.Context, flat, id string) error {
status, payload, err := s.do(ctx, http.MethodDelete,
s.filesURL+"/"+url.PathEscape(id)+"?supportsAllDrives=true", "", nil)
if err != nil {
return err
}
s.forgetID(flat)
switch status {
case http.StatusOK, http.StatusNoContent, http.StatusNotFound:
return nil
default:
return errors.New("Drive deletion of ", flat, " answered ", status, ": ", string(payload))
}
}
func (s *driveStorage) Delete(ctx context.Context, name string) error {
flat := flatten(name)
if id, err := s.resolveID(ctx, flat); err == nil {
if err := s.deleteID(ctx, flat, id); err != nil {
return err
}
} else if err != errNotFound {
return err
}
children, err := s.query(ctx, "name contains '"+quoteDriveValue(flat+flatSeparator)+"'")
if err != nil {
return err
}
for childName, file := range children {
if err := s.deleteID(ctx, childName, file.id); err != nil {
return err
}
}
return nil
}
func (s *driveStorage) List(ctx context.Context, prefix string) ([]Entry, error) {
flat := flatten(prefix) + flatSeparator
found, err := s.query(ctx, "name contains '"+quoteDriveValue(flat)+"'")
if err != nil {
return nil, err
}
seen := make(map[string]bool, len(found))
entries := make([]Entry, 0, len(found))
for name, file := range found {
rest := strings.TrimPrefix(name, flat)
if rest == "" {
continue
}
direct := true
if cut := strings.Index(rest, flatSeparator); cut >= 0 {
rest = rest[:cut]
direct = false
}
if seen[rest] {
continue
}
seen[rest] = true
entry := Entry{Name: rest}
if direct && file.description != "" {
if data, err := base64.StdEncoding.DecodeString(file.description); err == nil {
entry.Inline = data
}
}
entries = append(entries, entry)
}
return entries, nil
}
func (s *driveStorage) Close() error {
return nil
}
@@ -0,0 +1,345 @@
package xdrive
import (
"bytes"
"context"
"crypto/rand"
"encoding/json"
"fmt"
"io"
"os"
"strconv"
"sync"
"testing"
"time"
"github.com/xtls/xray-core/transport/internet"
)
const liveSecretsEnv = "XRAY_XDRIVE_DRIVE_SECRETS"
func liveDriveConfig(t *testing.T) *Config {
t.Helper()
path := os.Getenv(liveSecretsEnv)
if path == "" {
t.Skipf("set %s to run this test", liveSecretsEnv)
}
payload, err := os.ReadFile(path)
if err != nil {
t.Fatalf("reading %s: %v", path, err)
}
var secrets struct {
Folder string `json:"folder"`
ClientID string `json:"client_id"`
ClientSecret string `json:"client_secret"`
RefreshToken string `json:"refresh_token"`
}
if err := json.Unmarshal(payload, &secrets); err != nil {
t.Fatalf("parsing %s: %v", path, err)
}
config := &Config{
RemoteFolder: secrets.Folder,
Service: "Google Drive",
Secrets: []string{secrets.ClientID, secrets.ClientSecret, secrets.RefreshToken},
SegmentBytes: 256 * 1024,
FlushIntervalMs: 100,
PollIntervalMs: 500,
MaxPollIntervalMs: 2000,
SessionTtlSeconds: 120,
}
if raw := os.Getenv("XRAY_XDRIVE_LIVE_SEGMENT"); raw != "" {
config.SegmentBytes = uint32(envInt(t, "XRAY_XDRIVE_LIVE_SEGMENT"))
}
if raw := os.Getenv("XRAY_XDRIVE_LIVE_CONCURRENCY"); raw != "" {
config.Concurrency = uint32(envInt(t, "XRAY_XDRIVE_LIVE_CONCURRENCY"))
}
storage, err := newDriveStorage(nil, config)
if err != nil {
t.Fatalf("newDriveStorage: %v", err)
}
defer storage.Close()
for _, dir := range []string{sessionsDir, streamsDir} {
if err := storage.Delete(context.Background(), dir); err != nil {
t.Fatalf("clearing %s: %v", dir, err)
}
}
return config
}
func envInt(t *testing.T, name string) int {
t.Helper()
parsed, err := strconv.Atoi(os.Getenv(name))
if err != nil {
t.Fatalf("%s: %v", name, err)
}
return parsed
}
func TestLiveDriveStorage(t *testing.T) {
config := liveDriveConfig(t)
storage, err := newDriveStorage(nil, config)
if err != nil {
t.Fatalf("newDriveStorage: %v", err)
}
defer storage.Close()
ctx := context.Background()
name := "streams/livetest/c2s/000000000.seg"
payload := []byte("xdrive over a real remote storage service")
defer storage.Delete(ctx, "streams/livetest")
start := time.Now()
if err := storage.Put(ctx, name, payload); err != nil {
t.Fatalf("Put: %v", err)
}
t.Logf("Put took %v", time.Since(start))
start = time.Now()
names, err := storage.List(ctx, "streams/livetest/c2s")
if err != nil {
t.Fatalf("List: %v", err)
}
t.Logf("List took %v", time.Since(start))
if len(names) != 1 || names[0].Name != "000000000.seg" {
t.Fatalf("List returned %v, want one segment", names)
}
start = time.Now()
got, err := storage.Get(ctx, name)
if err != nil {
t.Fatalf("Get: %v", err)
}
t.Logf("Get took %v", time.Since(start))
if !bytes.Equal(got, payload) {
t.Fatalf("Get returned %q, want %q", got, payload)
}
if _, err := storage.Get(ctx, "streams/livetest/c2s/000000009.seg"); err != errNotFound {
t.Fatalf("Get returned %v, want errNotFound", err)
}
if err := storage.Delete(ctx, "streams/livetest"); err != nil {
t.Fatalf("Delete: %v", err)
}
names, err = storage.List(ctx, "streams/livetest/c2s")
if err != nil {
t.Fatalf("List after delete: %v", err)
}
if len(names) != 0 {
t.Fatalf("List after delete returned %v", names)
}
}
func TestLiveDriveTransport(t *testing.T) {
config := liveDriveConfig(t)
streamSettings := &internet.MemoryStreamConfig{
ProtocolName: protocolName,
ProtocolSettings: config,
}
client, server, cleanup := pairWith(t, streamSettings)
defer cleanup()
start := time.Now()
if _, err := client.Write([]byte("ping")); err != nil {
t.Fatalf("client write: %v", err)
}
expectRead(t, server, "ping")
t.Logf("client to server round took %v", time.Since(start))
start = time.Now()
if _, err := server.Write([]byte("pong")); err != nil {
t.Fatalf("server write: %v", err)
}
expectRead(t, client, "pong")
t.Logf("server to client round took %v", time.Since(start))
size := 400000
if raw := os.Getenv("XRAY_XDRIVE_LIVE_BYTES"); raw != "" {
parsed, err := strconv.Atoi(raw)
if err != nil {
t.Fatalf("XRAY_XDRIVE_LIVE_BYTES: %v", err)
}
size = parsed
}
payload := make([]byte, size)
if _, err := rand.Read(payload); err != nil {
t.Fatalf("rand: %v", err)
}
start = time.Now()
go func() {
client.Write(payload)
}()
if err := server.SetReadDeadline(time.Now().Add(5 * time.Minute)); 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)
}
elapsed := time.Since(start)
if !bytes.Equal(got, payload) {
t.Fatal("payload mismatch")
}
t.Logf("%d bytes took %v (%.1f KiB/s)", len(payload), elapsed,
float64(len(payload))/1024/elapsed.Seconds())
}
func TestLiveDriveParallelPut(t *testing.T) {
config := liveDriveConfig(t)
storage, err := newDriveStorage(nil, config)
if err != nil {
t.Fatalf("newDriveStorage: %v", err)
}
defer storage.Close()
ctx := context.Background()
defer storage.Delete(ctx, "streams/benchtest")
chunk := make([]byte, 256*1024)
if _, err := rand.Read(chunk); err != nil {
t.Fatalf("rand: %v", err)
}
if err := storage.Put(ctx, "streams/benchtest/warmup", chunk); err != nil {
t.Fatalf("warmup: %v", err)
}
start := time.Now()
for i := 0; i < 4; i++ {
if err := storage.Put(ctx, fmt.Sprintf("streams/benchtest/seq%d", i), chunk); err != nil {
t.Fatalf("sequential put: %v", err)
}
}
sequential := time.Since(start)
t.Logf("4 sequential puts of 256 KiB: %v (%.1f KiB/s)",
sequential, float64(4*len(chunk))/1024/sequential.Seconds())
start = time.Now()
var wg sync.WaitGroup
failures := make([]error, 8)
for i := 0; i < 8; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
failures[i] = storage.Put(ctx, fmt.Sprintf("streams/benchtest/par%d", i), chunk)
}(i)
}
wg.Wait()
parallel := time.Since(start)
for _, err := range failures {
if err != nil {
t.Fatalf("parallel put: %v", err)
}
}
t.Logf("8 parallel puts of 256 KiB: %v (%.1f KiB/s)",
parallel, float64(8*len(chunk))/1024/parallel.Seconds())
start = time.Now()
names, err := storage.List(ctx, "streams/benchtest")
if err != nil {
t.Fatalf("List: %v", err)
}
t.Logf("List of %d objects took %v", len(names), time.Since(start))
start = time.Now()
wg = sync.WaitGroup{}
for i := 0; i < 8; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
storage.Get(ctx, fmt.Sprintf("streams/benchtest/par%d", i))
}(i)
}
wg.Wait()
download := time.Since(start)
t.Logf("8 parallel gets of 256 KiB: %v (%.1f KiB/s)",
download, float64(8*len(chunk))/1024/download.Seconds())
}
func TestLiveDriveSegmentSweep(t *testing.T) {
config := liveDriveConfig(t)
storage, err := newDriveStorage(nil, config)
if err != nil {
t.Fatalf("newDriveStorage: %v", err)
}
defer storage.Close()
ctx := context.Background()
defer storage.Delete(ctx, "streams/sweeptest")
const total = 1024 * 1024
for _, size := range []int{64 * 1024, 128 * 1024, 256 * 1024, 512 * 1024} {
chunk := make([]byte, size)
if _, err := rand.Read(chunk); err != nil {
t.Fatalf("rand: %v", err)
}
count := total / size
start := time.Now()
var wg sync.WaitGroup
for i := 0; i < count; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
storage.Put(ctx, fmt.Sprintf("streams/sweeptest/s%d-%d", size, i), chunk)
}(i)
}
wg.Wait()
elapsed := time.Since(start)
t.Logf("%4d KiB x %2d = 1 MiB in %8v -> %6.1f KiB/s",
size/1024, count, elapsed.Round(time.Millisecond),
float64(total)/1024/elapsed.Seconds())
}
}
func TestLiveDriveListLag(t *testing.T) {
config := liveDriveConfig(t)
storage, err := newDriveStorage(nil, config)
if err != nil {
t.Fatalf("newDriveStorage: %v", err)
}
defer storage.Close()
ctx := context.Background()
defer storage.Delete(ctx, "streams/lagtest")
const rounds = 6
var worst time.Duration
for i := 0; i < rounds; i++ {
name := fmt.Sprintf("streams/lagtest/round%d/000000000.seg", i)
if err := storage.Put(ctx, name, []byte("probe")); err != nil {
t.Fatalf("Put: %v", err)
}
start := time.Now()
var lag time.Duration
for {
names, err := storage.List(ctx, fmt.Sprintf("streams/lagtest/round%d", i))
if err != nil {
t.Fatalf("List: %v", err)
}
if len(names) == 1 {
lag = time.Since(start)
break
}
if time.Since(start) > 30*time.Second {
t.Fatalf("round %d: the object never showed up in a listing", i)
}
}
if lag > worst {
worst = lag
}
t.Logf("round %d: the object became listable after %v", i, lag.Round(time.Millisecond))
}
t.Logf("worst listing lag: %v", worst.Round(time.Millisecond))
}
+690
View File
@@ -0,0 +1,690 @@
package xdrive
import (
"context"
"encoding/json"
"fmt"
"io"
"mime"
"mime/multipart"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"sync"
"testing"
"time"
"github.com/xtls/xray-core/common/net"
"github.com/xtls/xray-core/transport/internet"
)
type fakeFile struct {
id string
name string
data []byte
description string
}
type fakeDrive struct {
server *httptest.Server
mu sync.Mutex
files map[string]*fakeFile
nextID int
failOnce map[string]bool
failStatus map[string]int
failBody map[string]string
tokens int
hosts map[string]bool
}
func newFakeDrive(t *testing.T) *fakeDrive {
t.Helper()
drive := &fakeDrive{
files: make(map[string]*fakeFile),
failOnce: make(map[string]bool),
failStatus: make(map[string]int),
failBody: make(map[string]string),
hosts: make(map[string]bool),
}
mux := http.NewServeMux()
mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.NotFound(w, r)
}))
record := func(next http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
drive.mu.Lock()
drive.hosts[r.Host] = true
drive.mu.Unlock()
next(w, r)
}
}
mux.HandleFunc("/token", record(drive.handleToken))
mux.HandleFunc("/upload", record(drive.handleUpload))
mux.HandleFunc("/files", record(drive.handleFiles))
mux.HandleFunc("/files/", record(drive.handleFile))
drive.server = httptest.NewServer(mux)
resetSharedStorage()
previous := []string{driveTokenURL, driveFilesURL, driveUploadURL}
previousBackoff := driveInitialBackoff
driveTokenURL = drive.server.URL + "/token"
driveFilesURL = drive.server.URL + "/files"
driveUploadURL = drive.server.URL + "/upload"
driveInitialBackoff = 5 * time.Millisecond
t.Cleanup(func() {
driveTokenURL, driveFilesURL, driveUploadURL = previous[0], previous[1], previous[2]
driveInitialBackoff = previousBackoff
drive.server.Close()
resetSharedStorage()
})
return drive
}
func (d *fakeDrive) handleToken(w http.ResponseWriter, r *http.Request) {
d.mu.Lock()
d.tokens++
d.mu.Unlock()
json.NewEncoder(w).Encode(map[string]interface{}{
"access_token": "fake-token",
"expires_in": 3600,
})
}
func (d *fakeDrive) handleUpload(w http.ResponseWriter, r *http.Request) {
if d.shouldFail("upload") {
if status, body := d.failure("upload"); status != 0 {
w.WriteHeader(status)
w.Write([]byte(body))
return
}
w.WriteHeader(http.StatusTooManyRequests)
return
}
_, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
if err != nil {
w.WriteHeader(http.StatusBadRequest)
return
}
reader := multipart.NewReader(r.Body, params["boundary"])
metaPart, err := reader.NextPart()
if err != nil {
w.WriteHeader(http.StatusBadRequest)
return
}
var metadata struct {
Name string `json:"name"`
}
if err := json.NewDecoder(metaPart).Decode(&metadata); err != nil {
w.WriteHeader(http.StatusBadRequest)
return
}
var data []byte
if dataPart, err := reader.NextPart(); err == nil {
data, _ = io.ReadAll(dataPart)
}
d.mu.Lock()
d.nextID++
id := fmt.Sprintf("id-%d", d.nextID)
d.files[id] = &fakeFile{id: id, name: metadata.Name, data: data}
d.mu.Unlock()
json.NewEncoder(w).Encode(map[string]string{"id": id, "name": metadata.Name})
}
func (d *fakeDrive) handleFiles(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodPost {
d.handleCreate(w, r)
return
}
d.handleList(w, r)
}
func (d *fakeDrive) handleCreate(w http.ResponseWriter, r *http.Request) {
if d.shouldFail("upload") {
if status, body := d.failure("upload"); status != 0 {
w.WriteHeader(status)
w.Write([]byte(body))
return
}
w.WriteHeader(http.StatusTooManyRequests)
return
}
var meta struct {
Name string `json:"name"`
Description string `json:"description"`
}
if err := json.NewDecoder(r.Body).Decode(&meta); err != nil {
w.WriteHeader(http.StatusBadRequest)
return
}
d.mu.Lock()
d.nextID++
id := fmt.Sprintf("id-%d", d.nextID)
d.files[id] = &fakeFile{id: id, name: meta.Name, description: meta.Description}
d.mu.Unlock()
json.NewEncoder(w).Encode(map[string]string{"id": id, "name": meta.Name})
}
func (d *fakeDrive) handleList(w http.ResponseWriter, r *http.Request) {
if d.shouldFail("list") {
w.WriteHeader(http.StatusServiceUnavailable)
return
}
query := r.URL.Query().Get("q")
exact, prefix := parseFakeQuery(query)
type entry struct {
ID string `json:"id"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
}
result := struct {
Files []entry `json:"files"`
}{}
d.mu.Lock()
for _, file := range d.files {
match := false
switch {
case exact != "":
match = file.name == exact
case prefix != "":
match = strings.HasPrefix(file.name, prefix)
}
if match {
result.Files = append(result.Files, entry{
ID: file.id, Name: file.name, Description: file.description,
})
}
}
d.mu.Unlock()
json.NewEncoder(w).Encode(result)
}
func (d *fakeDrive) handleFile(w http.ResponseWriter, r *http.Request) {
id := strings.TrimPrefix(r.URL.Path, "/files/")
d.mu.Lock()
file, ok := d.files[id]
if ok && r.Method == http.MethodDelete {
delete(d.files, id)
}
d.mu.Unlock()
if !ok {
w.WriteHeader(http.StatusNotFound)
return
}
if r.Method == http.MethodDelete {
w.WriteHeader(http.StatusNoContent)
return
}
if strings.Contains(r.URL.RawQuery, "fields=description") {
json.NewEncoder(w).Encode(map[string]string{"description": file.description})
return
}
w.Write(file.data)
}
func (d *fakeDrive) shouldFail(kind string) bool {
d.mu.Lock()
defer d.mu.Unlock()
if d.failOnce[kind] {
d.failOnce[kind] = false
return true
}
return false
}
func (d *fakeDrive) failOnceWith(kind string, status int, body string) {
d.mu.Lock()
d.failOnce[kind] = true
d.failStatus[kind] = status
d.failBody[kind] = body
d.mu.Unlock()
}
func (d *fakeDrive) failure(kind string) (int, string) {
d.mu.Lock()
defer d.mu.Unlock()
if status, ok := d.failStatus[kind]; ok {
return status, d.failBody[kind]
}
return 0, ""
}
func (d *fakeDrive) failNext(kind string) {
d.mu.Lock()
d.failOnce[kind] = true
d.mu.Unlock()
}
func (d *fakeDrive) seenHosts() []string {
d.mu.Lock()
defer d.mu.Unlock()
out := make([]string, 0, len(d.hosts))
for h := range d.hosts {
out = append(out, h)
}
return out
}
func (d *fakeDrive) count() int {
d.mu.Lock()
defer d.mu.Unlock()
return len(d.files)
}
func parseFakeQuery(query string) (exact, prefix string) {
if value, ok := cutQuoted(query, "name = '"); ok {
return value, ""
}
if value, ok := cutQuoted(query, "name contains '"); ok {
return "", value
}
return "", ""
}
func cutQuoted(query, marker string) (string, bool) {
start := strings.Index(query, marker)
if start < 0 {
return "", false
}
rest := query[start+len(marker):]
end := strings.Index(rest, "'")
if end < 0 {
return "", false
}
return rest[:end], true
}
func driveSettings() *internet.MemoryStreamConfig {
return &internet.MemoryStreamConfig{
ProtocolName: protocolName,
ProtocolSettings: &Config{
RemoteFolder: "folder-id",
Service: "Google Drive",
Secrets: []string{"client", "secret", "refresh"},
FlushIntervalMs: 5,
PollIntervalMs: 5,
MaxPollIntervalMs: 20,
SessionTtlSeconds: 1,
},
}
}
func newDriveBackend(t *testing.T) *driveStorage {
t.Helper()
storage, err := newDriveStorage(driveSettings(), driveSettings().ProtocolSettings.(*Config))
if err != nil {
t.Fatalf("newDriveStorage: %v", err)
}
return storage
}
func TestDriveSecrets(t *testing.T) {
if _, err := newDriveStorage(nil, &Config{RemoteFolder: "f", Secrets: []string{"a", "b"}}); err == nil {
t.Fatal("accepted two secrets")
}
if _, err := newDriveStorage(nil, &Config{Secrets: []string{"a", "b", "c"}}); err == nil {
t.Fatal("accepted an empty remoteFolder")
}
if _, err := newDriveStorage(nil, &Config{RemoteFolder: "f", Secrets: []string{"a", "", "c"}}); err == nil {
t.Fatal("accepted an empty secret")
}
}
func TestDriveRoundTrip(t *testing.T) {
drive := newFakeDrive(t)
storage := newDriveBackend(t)
ctx := context.Background()
if err := storage.Put(ctx, "streams/abc/c2s/000000000.seg", []byte("hello")); err != nil {
t.Fatalf("Put: %v", err)
}
data, err := storage.Get(ctx, "streams/abc/c2s/000000000.seg")
if err != nil {
t.Fatalf("Get: %v", err)
}
if string(data) != "hello" {
t.Fatalf("Get returned %q, want %q", data, "hello")
}
names, err := storage.List(ctx, "streams/abc/c2s")
if err != nil {
t.Fatalf("List: %v", err)
}
if len(names) != 1 || names[0].Name != "000000000.seg" {
t.Fatalf("List returned %v, want 1 segment", names)
}
if _, err := storage.Get(ctx, "streams/abc/c2s/000000009.seg"); err != errNotFound {
t.Fatalf("Get returned %v, want errNotFound", err)
}
if err := storage.Delete(ctx, "streams/abc/c2s/000000000.seg"); err != nil {
t.Fatalf("Delete: %v", err)
}
if drive.count() != 0 {
t.Fatalf("fake drive still holds %d files", drive.count())
}
}
func TestDriveListChildren(t *testing.T) {
newFakeDrive(t)
storage := newDriveBackend(t)
ctx := context.Background()
for _, name := range []string{
"streams/one/c2s/000000000.seg",
"streams/one/s2c/000000000.seg",
"streams/two/c2s/000000000.seg",
} {
if err := storage.Put(ctx, name, []byte("x")); err != nil {
t.Fatalf("Put %s: %v", name, err)
}
}
names, err := storage.List(ctx, "streams")
if err != nil {
t.Fatalf("List: %v", err)
}
if len(names) != 2 {
t.Fatalf("List returned %v, want 2 sessions", names)
}
}
func TestDriveDeleteSession(t *testing.T) {
drive := newFakeDrive(t)
storage := newDriveBackend(t)
ctx := context.Background()
for _, name := range []string{
"streams/one/c2s/000000000.seg",
"streams/one/c2s/000000001.end",
"streams/one/s2c/000000000.seg",
"streams/two/c2s/000000000.seg",
} {
if err := storage.Put(ctx, name, []byte("x")); err != nil {
t.Fatalf("Put %s: %v", name, err)
}
}
if err := storage.Delete(ctx, "streams/one"); err != nil {
t.Fatalf("Delete: %v", err)
}
if drive.count() != 1 {
t.Fatalf("fake drive holds %d files, want 1", drive.count())
}
}
func TestDriveRetry(t *testing.T) {
drive := newFakeDrive(t)
storage := newDriveBackend(t)
ctx := context.Background()
drive.failNext("upload")
if err := storage.Put(ctx, "sessions/abc", nil); err != nil {
t.Fatalf("Put did not survive a 429: %v", err)
}
drive.failNext("list")
if _, err := storage.List(ctx, "sessions"); err != nil {
t.Fatalf("List did not survive a 503: %v", err)
}
}
func TestDriveTokenCache(t *testing.T) {
drive := newFakeDrive(t)
storage := newDriveBackend(t)
ctx := context.Background()
for i := 0; i < 5; i++ {
if err := storage.Put(ctx, fmt.Sprintf("sessions/s%d", i), nil); err != nil {
t.Fatalf("Put: %v", err)
}
}
if drive.tokens != 1 {
t.Fatalf("token endpoint hit %d times, want 1", drive.tokens)
}
}
func TestDriveTransport(t *testing.T) {
newFakeDrive(t)
client, server, cleanup := pairWith(t, driveSettings())
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 TestDriveLargeTransfer(t *testing.T) {
newFakeDrive(t)
client, server, cleanup := pairWith(t, driveSettings())
defer cleanup()
payload := make([]byte, 300000)
for i := range payload {
payload[i] = byte(i % 251)
}
go func() {
client.Write(payload)
}()
if err := server.SetReadDeadline(time.Now().Add(60 * 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)
}
for i := range got {
if got[i] != payload[i] {
t.Fatalf("payload mismatch at byte %d", i)
}
}
}
func TestRateLimited(t *testing.T) {
limited := []string{
`{"error":{"code":403,"errors":[{"reason":"userRateLimitExceeded"}]}}`,
`{"error":{"code":403,"errors":[{"reason":"rateLimitExceeded"}]}}`,
`{"error":{"code":403,"errors":[{"reason":"sharingRateLimitExceeded"}]}}`,
`{"error":{"status":"RESOURCE_EXHAUSTED"}}`,
}
for _, payload := range limited {
if !rateLimited([]byte(payload)) {
t.Fatalf("rateLimited missed %s", payload)
}
}
permanent := []string{
`{"error":{"code":403,"errors":[{"reason":"insufficientFilePermissions"}]}}`,
`{"error":{"code":403,"errors":[{"reason":"storageQuotaExceeded"}]}}`,
`not json at all`,
}
for _, payload := range permanent {
if rateLimited([]byte(payload)) {
t.Fatalf("rateLimited treated %s as temporary", payload)
}
}
}
func TestDriveRetryRateLimit(t *testing.T) {
drive := newFakeDrive(t)
storage := newDriveBackend(t)
drive.failOnceWith("upload", http.StatusForbidden,
`{"error":{"code":403,"errors":[{"reason":"userRateLimitExceeded"}]}}`)
if err := storage.Put(context.Background(), "sessions/abc", nil); err != nil {
t.Fatalf("Put did not survive a 403: %v", err)
}
}
func TestDriveSharedClient(t *testing.T) {
newFakeDrive(t)
first, err := newStorage(driveSettings())
if err != nil {
t.Fatalf("newStorage: %v", err)
}
second, err := newStorage(driveSettings())
if err != nil {
t.Fatalf("newStorage: %v", err)
}
if first != second {
t.Fatal("same settings did not share one storage")
}
}
func TestDriveInlineListing(t *testing.T) {
drive := newFakeDrive(t)
storage := newDriveBackend(t)
ctx := context.Background()
payload := []byte("small enough to ride along with the listing")
if err := storage.Put(ctx, "streams/abc/c2s/000000000.seg", payload); err != nil {
t.Fatalf("Put: %v", err)
}
entries, err := storage.List(ctx, "streams/abc/c2s")
if err != nil {
t.Fatalf("List: %v", err)
}
if len(entries) != 1 {
t.Fatalf("List returned %v, want 1 entry", entries)
}
if string(entries[0].Inline) != string(payload) {
t.Fatalf("listing carried %q, want %q", entries[0].Inline, payload)
}
d := drive
d.mu.Lock()
var stored *fakeFile
for _, f := range d.files {
stored = f
}
d.mu.Unlock()
if len(stored.data) != 0 {
t.Fatal("small payload was uploaded as content")
}
}
func TestDriveLargeContent(t *testing.T) {
drive := newFakeDrive(t)
storage := newDriveBackend(t)
ctx := context.Background()
payload := make([]byte, driveInlineLimit+1)
for i := range payload {
payload[i] = byte(i)
}
if err := storage.Put(ctx, "streams/abc/c2s/000000000.seg", payload); err != nil {
t.Fatalf("Put: %v", err)
}
entries, err := storage.List(ctx, "streams/abc/c2s")
if err != nil {
t.Fatalf("List: %v", err)
}
if len(entries) != 1 || entries[0].Inline != nil {
t.Fatalf("large payload was inlined, got %v", entries)
}
got, err := storage.Get(ctx, "streams/abc/c2s/000000000.seg")
if err != nil {
t.Fatalf("Get: %v", err)
}
if len(got) != len(payload) {
t.Fatalf("Get returned %d bytes, want %d", len(got), len(payload))
}
_ = drive
}
func TestDriveInlineGet(t *testing.T) {
newFakeDrive(t)
storage := newDriveBackend(t)
ctx := context.Background()
payload := []byte("only in the description")
if err := storage.Put(ctx, "sessions/abc", payload); err != nil {
t.Fatalf("Put: %v", err)
}
got, err := storage.Get(ctx, "sessions/abc")
if err != nil {
t.Fatalf("Get: %v", err)
}
if string(got) != string(payload) {
t.Fatalf("Get returned %q, want %q", got, payload)
}
}
func TestDriveFronting(t *testing.T) {
drive := newFakeDrive(t)
fake, err := url.Parse(drive.server.URL)
if err != nil {
t.Fatalf("parse: %v", err)
}
port, err := net.PortFromString(fake.Port())
if err != nil {
t.Fatalf("port: %v", err)
}
driveTokenURL = "http://www.googleapis.com/token"
driveFilesURL = "http://www.googleapis.com/files"
driveUploadURL = "http://www.googleapis.com/upload"
settings := driveSettings()
settings.Destination = &net.Destination{
Address: net.ParseAddress(fake.Hostname()),
Port: port,
Network: net.Network_TCP,
}
storage, err := newDriveStorage(settings, settings.ProtocolSettings.(*Config))
if err != nil {
t.Fatalf("newDriveStorage: %v", err)
}
if err := storage.Put(context.Background(), "sessions/fronted", []byte("x")); err != nil {
t.Fatalf("Put: %v", err)
}
if drive.count() != 1 {
t.Fatalf("fake drive holds %d files, want 1", drive.count())
}
for _, host := range drive.seenHosts() {
if host != "www.googleapis.com" {
t.Fatalf("inner host was %q, want www.googleapis.com regardless of address", host)
}
}
}
+39
View File
@@ -0,0 +1,39 @@
package xdrive
import (
"encoding/json"
"strings"
)
func jsonWalk(payload []byte, path string) interface{} {
var root interface{}
if json.Unmarshal(payload, &root) != nil {
return nil
}
node := root
for _, key := range strings.Split(path, ".") {
obj, ok := node.(map[string]interface{})
if !ok {
return nil
}
node, ok = obj[key]
if !ok {
return nil
}
}
return node
}
func jsonString(payload []byte, path string) string {
if s, ok := jsonWalk(payload, path).(string); ok {
return s
}
return ""
}
func jsonNumber(payload []byte, path string) int64 {
if f, ok := jsonWalk(payload, path).(float64); ok {
return int64(f)
}
return 0
}
+120
View File
@@ -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
}
+69
View File
@@ -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
}
+90
View File
@@ -0,0 +1,90 @@
package xdrive
import (
"context"
"crypto/sha256"
"encoding/hex"
"strings"
"sync"
"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 sharedStorage(streamSettings, config, func() (Storage, error) {
return newDriveStorage(streamSettings, config)
})
case "template":
return sharedStorage(streamSettings, config, func() (Storage, error) {
return newTemplateStorage(streamSettings, config)
})
default:
return nil, errors.New("unsupported service: ", config.Service)
}
}
var (
sharedMu sync.Mutex
shared = make(map[string]Storage)
)
func shareKey(streamSettings *internet.MemoryStreamConfig, config *Config) string {
parts := []string{config.Service, config.RemoteFolder}
parts = append(parts, config.Secrets...)
if streamSettings != nil {
parts = append(parts, streamSettings.SecurityType)
if streamSettings.Destination != nil {
parts = append(parts, streamSettings.Destination.NetAddr())
}
}
sum := sha256.Sum256([]byte(strings.Join(parts, "\x00")))
return hex.EncodeToString(sum[:])
}
func resetSharedStorage() {
sharedMu.Lock()
shared = make(map[string]Storage)
sharedMu.Unlock()
}
func sharedStorage(streamSettings *internet.MemoryStreamConfig, config *Config, build func() (Storage, error)) (Storage, error) {
key := shareKey(streamSettings, config)
sharedMu.Lock()
defer sharedMu.Unlock()
if storage, ok := shared[key]; ok {
return storage, nil
}
storage, err := build()
if err != nil {
return nil, err
}
shared[key] = storage
return storage, nil
}
+445
View File
@@ -0,0 +1,445 @@
package xdrive
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"io"
"net/http"
"regexp"
"strings"
"sync"
"time"
"github.com/xtls/xray-core/common/errors"
"github.com/xtls/xray-core/transport/internet"
)
type authTemplate struct {
Type string `json:"type"`
Header map[string]string `json:"header"`
Username string `json:"username"`
Password string `json:"password"`
TokenURL string `json:"tokenUrl"`
Form map[string]string `json:"form"`
TokenPath string `json:"tokenPath"`
ExpiryPath string `json:"expiryPath"`
}
type opTemplate struct {
Method string `json:"method"`
URL string `json:"url"`
Headers map[string]string `json:"headers"`
Body string `json:"body"`
NamesRegex string `json:"namesRegex"`
}
type retryTemplate struct {
Status []int `json:"status"`
RateReason string `json:"rateReason"`
}
type storageTemplate struct {
Flatten bool `json:"flatten"`
Concurrency int `json:"concurrency"`
Auth authTemplate `json:"auth"`
Put opTemplate `json:"put"`
Get opTemplate `json:"get"`
Delete opTemplate `json:"delete"`
List opTemplate `json:"list"`
Retry retryTemplate `json:"retry"`
names *regexp.Regexp
}
type templateStorage struct {
tmpl *storageTemplate
client *http.Client
folder string
secrets []string
inflight chan struct{}
tokenMu sync.Mutex
token string
tokenExpiry time.Time
}
func newTemplateStorage(streamSettings *internet.MemoryStreamConfig, config *Config) (*templateStorage, error) {
tmpl := &storageTemplate{}
if err := json.Unmarshal([]byte(config.Template), tmpl); err != nil {
return nil, errors.New("invalid template").Base(err)
}
if tmpl.Put.URL == "" || tmpl.Get.URL == "" || tmpl.List.URL == "" || tmpl.Delete.URL == "" {
return nil, errors.New("template needs put, get, list and delete operations")
}
if tmpl.List.NamesRegex == "" {
return nil, errors.New("template list needs a namesRegex")
}
re, err := regexp.Compile(tmpl.List.NamesRegex)
if err != nil {
return nil, errors.New("bad namesRegex").Base(err)
}
if re.NumSubexp() < 1 {
return nil, errors.New("namesRegex needs one capture group")
}
tmpl.names = re
conc := tmpl.Concurrency
if conc <= 0 {
conc = driveMaxInflight
}
if conc > maxTemplateConcurrency {
conc = maxTemplateConcurrency
}
return &templateStorage{
tmpl: tmpl,
client: newServiceClient(streamSettings, driveTimeout, conc),
folder: config.RemoteFolder,
secrets: config.Secrets,
inflight: make(chan struct{}, conc),
}, nil
}
const maxTemplateConcurrency = 256
func (s *templateStorage) baseVars() map[string]string {
vars := map[string]string{"folder": s.folder}
for i, secret := range s.secrets {
vars["secret"+itoa(i)] = secret
}
return vars
}
func subst(tmpl string, vars map[string]string) string {
if tmpl == "" || !strings.ContainsRune(tmpl, '{') {
return tmpl
}
out := tmpl
for k, v := range vars {
out = strings.ReplaceAll(out, "{"+k+"}", v)
}
return out
}
func itoa(i int) string {
if i == 0 {
return "0"
}
var b [20]byte
pos := len(b)
for i > 0 {
pos--
b[pos] = byte('0' + i%10)
i /= 10
}
return string(b[pos:])
}
func (s *templateStorage) storedName(name string) string {
if s.tmpl.Flatten {
return flatten(name)
}
return name
}
func (s *templateStorage) retryable(status int, payload []byte) bool {
for _, code := range s.tmpl.Retry.Status {
if status == code {
return true
}
}
if status == http.StatusForbidden && s.tmpl.Retry.RateReason != "" {
if reason := jsonString(payload, s.tmpl.Retry.RateReason); reason != "" {
return true
}
}
return false
}
func (s *templateStorage) authHeaders(ctx context.Context, vars map[string]string) (map[string]string, error) {
switch s.tmpl.Auth.Type {
case "", "none":
return nil, nil
case "static", "oauth2":
if s.tmpl.Auth.Type == "oauth2" {
token, err := s.accessToken(ctx)
if err != nil {
return nil, err
}
vars["token"] = token
}
headers := make(map[string]string, len(s.tmpl.Auth.Header))
for k, v := range s.tmpl.Auth.Header {
headers[k] = subst(v, vars)
}
return headers, nil
case "basic":
user := subst(s.tmpl.Auth.Username, vars)
pass := subst(s.tmpl.Auth.Password, vars)
enc := base64.StdEncoding.EncodeToString([]byte(user + ":" + pass))
return map[string]string{"Authorization": "Basic " + enc}, nil
default:
return nil, errors.New("unsupported auth type: ", s.tmpl.Auth.Type)
}
}
func (s *templateStorage) accessToken(ctx context.Context) (string, error) {
s.tokenMu.Lock()
defer s.tokenMu.Unlock()
if s.token != "" && time.Now().Before(s.tokenExpiry) {
return s.token, nil
}
form := make(map[string]string, len(s.tmpl.Auth.Form))
vars := s.baseVars()
values := strings.Builder{}
first := true
for k, v := range s.tmpl.Auth.Form {
form[k] = subst(v, vars)
if !first {
values.WriteByte('&')
}
first = false
values.WriteString(k)
values.WriteByte('=')
values.WriteString(form[k])
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.tmpl.Auth.TokenURL,
strings.NewReader(values.String()))
if err != nil {
return "", errors.New("failed to build the token request").Base(err)
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
resp, err := s.client.Do(req)
if err != nil {
return "", errors.New("failed to fetch the token").Base(err)
}
defer resp.Body.Close()
payload, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
if err != nil {
return "", errors.New("failed to read the token response").Base(err)
}
if resp.StatusCode != http.StatusOK {
return "", errors.New("the token endpoint answered ", resp.StatusCode, ": ", string(payload))
}
path := s.tmpl.Auth.TokenPath
if path == "" {
path = "access_token"
}
token := jsonString(payload, path)
if token == "" {
return "", errors.New("the token response has no token at ", path)
}
lifetime := int64(3600)
if s.tmpl.Auth.ExpiryPath != "" {
if n := jsonNumber(payload, s.tmpl.Auth.ExpiryPath); n > 0 {
lifetime = n
}
}
if lifetime > 60 {
lifetime -= 60
}
s.token = token
s.tokenExpiry = time.Now().Add(time.Duration(lifetime) * time.Second)
return s.token, nil
}
func (s *templateStorage) invalidateToken() {
s.tokenMu.Lock()
s.token = ""
s.tokenMu.Unlock()
}
func (s *templateStorage) do(ctx context.Context, op *opTemplate, vars map[string]string, body []byte) (int, []byte, error) {
backoff := driveInitialBackoff
var lastErr error
for attempt := 0; attempt < driveMaxAttempts; attempt++ {
if attempt > 0 {
select {
case <-ctx.Done():
return 0, nil, ctx.Err()
case <-time.After(jitter(backoff)):
}
backoff *= 2
if backoff > driveMaxBackoff {
backoff = driveMaxBackoff
}
}
authHeaders, err := s.authHeaders(ctx, vars)
if err != nil {
lastErr = err
continue
}
method := op.Method
if method == "" {
method = http.MethodGet
}
var reader io.Reader
if body != nil {
reader = bytes.NewReader(body)
}
req, err := http.NewRequestWithContext(ctx, method, subst(op.URL, vars), reader)
if err != nil {
return 0, nil, errors.New("failed to build request").Base(err)
}
for k, v := range authHeaders {
req.Header.Set(k, v)
}
for k, v := range op.Headers {
req.Header.Set(k, subst(v, vars))
}
select {
case s.inflight <- struct{}{}:
case <-ctx.Done():
return 0, nil, ctx.Err()
}
resp, err := s.client.Do(req)
<-s.inflight
if err != nil {
if ctx.Err() != nil {
return 0, nil, ctx.Err()
}
lastErr = errors.New("request failed").Base(err)
continue
}
payload, err := io.ReadAll(resp.Body)
resp.Body.Close()
if err != nil {
if ctx.Err() != nil {
return 0, nil, ctx.Err()
}
lastErr = errors.New("failed to read response").Base(err)
continue
}
if resp.StatusCode == http.StatusUnauthorized && s.tmpl.Auth.Type == "oauth2" {
s.invalidateToken()
lastErr = errors.New("the service rejected the token")
continue
}
if s.retryable(resp.StatusCode, payload) {
lastErr = errors.New("the service answered ", resp.StatusCode)
continue
}
return resp.StatusCode, payload, nil
}
return 0, nil, lastErr
}
func (s *templateStorage) Put(ctx context.Context, name string, data []byte) error {
vars := s.baseVars()
vars["name"] = s.storedName(name)
body := data
if s.tmpl.Put.Body != "" {
vars["data"] = base64.StdEncoding.EncodeToString(data)
body = []byte(subst(s.tmpl.Put.Body, vars))
}
status, payload, err := s.do(ctx, &s.tmpl.Put, vars, body)
if err != nil {
return err
}
if status < 200 || status >= 300 {
return errors.New("put of ", name, " answered ", status, ": ", string(payload))
}
return nil
}
func (s *templateStorage) Get(ctx context.Context, name string) ([]byte, error) {
vars := s.baseVars()
vars["name"] = s.storedName(name)
status, payload, err := s.do(ctx, &s.tmpl.Get, vars, nil)
if err != nil {
return nil, err
}
switch {
case status >= 200 && status < 300:
return payload, nil
case status == http.StatusNotFound:
return nil, errNotFound
default:
return nil, errors.New("get of ", name, " answered ", status, ": ", string(payload))
}
}
func (s *templateStorage) Delete(ctx context.Context, name string) error {
vars := s.baseVars()
vars["name"] = s.storedName(name)
status, payload, err := s.do(ctx, &s.tmpl.Delete, vars, nil)
if err != nil {
return err
}
if status == http.StatusNotFound || (status >= 200 && status < 300) {
return nil
}
return errors.New("delete of ", name, " answered ", status, ": ", string(payload))
}
func (s *templateStorage) List(ctx context.Context, prefix string) ([]Entry, error) {
vars := s.baseVars()
flat := s.storedName(prefix)
vars["prefix"] = flat
status, payload, err := s.do(ctx, &s.tmpl.List, vars, nil)
if err != nil {
return nil, err
}
if status == http.StatusNotFound {
return nil, nil
}
if status < 200 || status >= 300 {
return nil, errors.New("list of ", prefix, " answered ", status, ": ", string(payload))
}
matches := s.tmpl.names.FindAllStringSubmatch(string(payload), -1)
if !s.tmpl.Flatten {
entries := make([]Entry, 0, len(matches))
for _, m := range matches {
entries = append(entries, Entry{Name: m[1]})
}
return entries, nil
}
want := flat + flatSeparator
seen := make(map[string]bool, len(matches))
entries := make([]Entry, 0, len(matches))
for _, m := range matches {
name := m[1]
if !strings.HasPrefix(name, want) {
continue
}
rest := strings.TrimPrefix(name, want)
if rest == "" {
continue
}
if cut := strings.Index(rest, flatSeparator); cut >= 0 {
rest = rest[:cut]
}
if seen[rest] {
continue
}
seen[rest] = true
entries = append(entries, Entry{Name: rest})
}
return entries, nil
}
func (s *templateStorage) Close() error {
return nil
}
+308
View File
@@ -0,0 +1,308 @@
package xdrive
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
"github.com/xtls/xray-core/transport/internet"
)
type fakeStore struct {
server *httptest.Server
mu sync.Mutex
objects map[string][]byte
needAuth string
sawAuth string
tokens int
}
func newFakeStore(t *testing.T) *fakeStore {
t.Helper()
store := &fakeStore{objects: make(map[string][]byte)}
store.server = httptest.NewServer(http.HandlerFunc(store.handle))
t.Cleanup(store.server.Close)
return store
}
func (s *fakeStore) handle(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/token" {
s.mu.Lock()
s.tokens++
s.mu.Unlock()
json.NewEncoder(w).Encode(map[string]interface{}{
"access_token": "tok-fake", "expires_in": 3600,
})
return
}
if auth := r.Header.Get("Authorization"); auth != "" {
s.mu.Lock()
s.sawAuth = auth
s.mu.Unlock()
}
if s.needAuth != "" && r.Header.Get("Authorization") != s.needAuth {
w.WriteHeader(http.StatusUnauthorized)
return
}
key := strings.TrimPrefix(r.URL.Path, "/folder/")
switch r.Method {
case "PROPFIND":
s.mu.Lock()
var b strings.Builder
for name := range s.objects {
fmt.Fprintf(&b, "<d:href>/folder/%s</d:href>\n", name)
}
s.mu.Unlock()
w.Write([]byte(b.String()))
case http.MethodPut:
body, _ := io.ReadAll(r.Body)
s.mu.Lock()
s.objects[key] = body
s.mu.Unlock()
w.WriteHeader(http.StatusCreated)
case http.MethodGet:
s.mu.Lock()
data, ok := s.objects[key]
s.mu.Unlock()
if !ok {
w.WriteHeader(http.StatusNotFound)
return
}
w.Write(data)
case http.MethodDelete:
s.mu.Lock()
delete(s.objects, key)
s.mu.Unlock()
w.WriteHeader(http.StatusNoContent)
default:
w.WriteHeader(http.StatusMethodNotAllowed)
}
}
func (s *fakeStore) count() int {
s.mu.Lock()
defer s.mu.Unlock()
return len(s.objects)
}
func (s *fakeStore) seenAuth() string {
s.mu.Lock()
defer s.mu.Unlock()
return s.sawAuth
}
func templateSettings(store *fakeStore, auth map[string]interface{}, secrets []string) *internet.MemoryStreamConfig {
base := store.server.URL
tmpl := map[string]interface{}{
"flatten": true,
"auth": auth,
"put": map[string]interface{}{"method": "PUT", "url": base + "/folder/{name}"},
"get": map[string]interface{}{"method": "GET", "url": base + "/folder/{name}"},
"delete": map[string]interface{}{"method": "DELETE", "url": base + "/folder/{name}"},
"list": map[string]interface{}{
"method": "PROPFIND",
"url": base + "/folder/",
"namesRegex": `<d:href>/folder/([^<]+)</d:href>`,
},
"retry": map[string]interface{}{"status": []int{429, 500, 502, 503}},
}
raw, _ := json.Marshal(tmpl)
return &internet.MemoryStreamConfig{
ProtocolName: protocolName,
ProtocolSettings: &Config{
RemoteFolder: "folder",
Service: "template",
Secrets: secrets,
Template: string(raw),
FlushIntervalMs: 5,
PollIntervalMs: 5,
MaxPollIntervalMs: 20,
SessionTtlSeconds: 5,
},
}
}
func newTemplateBackend(t *testing.T, store *fakeStore, auth map[string]interface{}, secrets []string) *templateStorage {
t.Helper()
settings := templateSettings(store, auth, secrets)
storage, err := newTemplateStorage(settings, settings.ProtocolSettings.(*Config))
if err != nil {
t.Fatalf("newTemplateStorage: %v", err)
}
return storage
}
func TestTemplateRoundTrip(t *testing.T) {
store := newFakeStore(t)
storage := newTemplateBackend(t, store, map[string]interface{}{"type": "none"}, nil)
ctx := context.Background()
if err := storage.Put(ctx, "streams/abc/c2s/000000000.seg", []byte("hello")); err != nil {
t.Fatalf("Put: %v", err)
}
data, err := storage.Get(ctx, "streams/abc/c2s/000000000.seg")
if err != nil {
t.Fatalf("Get: %v", err)
}
if string(data) != "hello" {
t.Fatalf("Get returned %q, want hello", data)
}
entries, err := storage.List(ctx, "streams/abc/c2s")
if err != nil {
t.Fatalf("List: %v", err)
}
if len(entries) != 1 || entries[0].Name != "000000000.seg" {
t.Fatalf("List returned %v, want one segment", entries)
}
if _, err := storage.Get(ctx, "streams/abc/c2s/000000009.seg"); err != errNotFound {
t.Fatalf("Get of a missing object returned %v, want errNotFound", err)
}
if err := storage.Delete(ctx, "streams/abc/c2s/000000000.seg"); err != nil {
t.Fatalf("Delete: %v", err)
}
if store.count() != 0 {
t.Fatalf("store still holds %d objects", store.count())
}
}
func TestTemplateListReturnsDirectChildren(t *testing.T) {
store := newFakeStore(t)
storage := newTemplateBackend(t, store, map[string]interface{}{"type": "none"}, nil)
ctx := context.Background()
for _, name := range []string{
"streams/one/c2s/000000000.seg",
"streams/one/s2c/000000000.seg",
"streams/two/c2s/000000000.seg",
} {
if err := storage.Put(ctx, name, []byte("x")); err != nil {
t.Fatalf("Put %s: %v", name, err)
}
}
entries, err := storage.List(ctx, "streams")
if err != nil {
t.Fatalf("List: %v", err)
}
if len(entries) != 2 {
t.Fatalf("List returned %v, want the two session ids", entries)
}
}
func TestTemplateBasicAuth(t *testing.T) {
store := newFakeStore(t)
store.needAuth = "Basic dXNlcjpwYXNz"
auth := map[string]interface{}{"type": "basic", "username": "{secret0}", "password": "{secret1}"}
storage := newTemplateBackend(t, store, auth, []string{"user", "pass"})
if err := storage.Put(context.Background(), "sessions/a", []byte("x")); err != nil {
t.Fatalf("Put with basic auth: %v", err)
}
if store.seenAuth() != "Basic dXNlcjpwYXNz" {
t.Fatalf("server saw auth %q", store.seenAuth())
}
}
func TestTemplateOAuth(t *testing.T) {
store := newFakeStore(t)
store.needAuth = "Bearer tok-fake"
auth := map[string]interface{}{
"type": "oauth2",
"tokenUrl": store.server.URL + "/token",
"form": map[string]interface{}{"grant_type": "refresh_token", "refresh_token": "{secret0}"},
"header": map[string]interface{}{"Authorization": "Bearer {token}"},
}
storage := newTemplateBackend(t, store, auth, []string{"refresh"})
ctx := context.Background()
for i := 0; i < 4; i++ {
if err := storage.Put(ctx, fmt.Sprintf("sessions/s%d", i), nil); err != nil {
t.Fatalf("Put: %v", err)
}
}
if store.tokens != 1 {
t.Fatalf("token endpoint was hit %d times, want 1", store.tokens)
}
}
func TestTemplateTransport(t *testing.T) {
store := newFakeStore(t)
settings := templateSettings(store, map[string]interface{}{"type": "none"}, nil)
client, server, cleanup := pairWith(t, settings)
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")
payload := make([]byte, 300000)
for i := range payload {
payload[i] = byte(i % 251)
}
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)
}
for i := range got {
if got[i] != payload[i] {
t.Fatalf("payload mismatch at byte %d", i)
}
}
}
func TestTemplateConcurrency(t *testing.T) {
store := newFakeStore(t)
settings := templateSettings(store, map[string]interface{}{"type": "none"}, nil)
var tmpl map[string]interface{}
cfg := settings.ProtocolSettings.(*Config)
if err := json.Unmarshal([]byte(cfg.Template), &tmpl); err != nil {
t.Fatalf("unmarshal: %v", err)
}
tmpl["concurrency"] = 4
raw, _ := json.Marshal(tmpl)
cfg.Template = string(raw)
storage, err := newTemplateStorage(settings, cfg)
if err != nil {
t.Fatalf("newTemplateStorage: %v", err)
}
if cap(storage.inflight) != 4 {
t.Fatalf("inflight cap is %d, want 4 from the template", cap(storage.inflight))
}
}
func TestTemplateConcurrencyDefault(t *testing.T) {
store := newFakeStore(t)
storage := newTemplateBackend(t, store, map[string]interface{}{"type": "none"}, nil)
if cap(storage.inflight) != driveMaxInflight {
t.Fatalf("default inflight cap is %d, want %d", cap(storage.inflight), driveMaxInflight)
}
}
+445
View File
@@ -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
}
+308
View File
@@ -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()
}
+583
View File
@@ -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)
}
}
@@ -0,0 +1,171 @@
package xdrive
import (
"bytes"
"context"
"crypto/rand"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"testing"
"time"
"github.com/xtls/xray-core/transport/internet"
)
const yandexBase = "https://webdav.yandex.ru"
func envUint(name string, def uint32) uint32 {
if v := os.Getenv(name); v != "" {
var n uint32
fmt.Sscanf(v, "%d", &n)
if n > 0 {
return n
}
}
return def
}
func liveYandexSettings(t *testing.T) (*internet.MemoryStreamConfig, string, func()) {
t.Helper()
user := os.Getenv("XDRIVE_YANDEX_USER")
pass := os.Getenv("XDRIVE_YANDEX_PASS")
if user == "" || pass == "" {
t.Skip("set XDRIVE_YANDEX_USER and XDRIVE_YANDEX_PASS to run this test")
}
folder := fmt.Sprintf("xdrive-live-%d", time.Now().UnixNano())
dav := func(method, path string) int {
req, _ := http.NewRequest(method, yandexBase+path, nil)
req.SetBasicAuth(user, pass)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatalf("%s %s: %v", method, path, err)
}
resp.Body.Close()
return resp.StatusCode
}
if code := dav("MKCOL", "/"+folder); code != 201 && code != 405 {
t.Fatalf("MKCOL answered %d", code)
}
tmpl := map[string]interface{}{
"flatten": true,
"auth": map[string]interface{}{"type": "basic", "username": "{secret0}", "password": "{secret1}"},
"put": map[string]interface{}{"method": "PUT", "url": yandexBase + "/{folder}/{name}"},
"get": map[string]interface{}{"method": "GET", "url": yandexBase + "/{folder}/{name}"},
"delete": map[string]interface{}{"method": "DELETE", "url": yandexBase + "/{folder}/{name}"},
"list": map[string]interface{}{
"method": "PROPFIND", "url": yandexBase + "/{folder}/",
"headers": map[string]interface{}{"Depth": "1"}, "namesRegex": `<d:href>[^<]*/([^/<]+)</d:href>`,
},
"retry": map[string]interface{}{"status": []int{429, 500, 502, 503}},
}
raw, _ := json.Marshal(tmpl)
settings := &internet.MemoryStreamConfig{
ProtocolName: protocolName,
ProtocolSettings: &Config{
RemoteFolder: folder,
Service: "template",
Secrets: []string{user, pass},
Template: string(raw),
SegmentBytes: 262144,
FlushIntervalMs: 100,
PollIntervalMs: 300,
MaxPollIntervalMs: 1500,
SessionTtlSeconds: 120,
Concurrency: envUint("XDRIVE_LIVE_CONCURRENCY", 8),
},
}
cleanup := func() { dav("DELETE", "/"+folder) }
return settings, folder, cleanup
}
func TestLiveYandexStorage(t *testing.T) {
settings, _, cleanup := liveYandexSettings(t)
defer cleanup()
storage, err := newTemplateStorage(settings, settings.ProtocolSettings.(*Config))
if err != nil {
t.Fatalf("newTemplateStorage: %v", err)
}
ctx := context.Background()
name := "streams/live/c2s/000000000.seg"
payload := []byte("xdrive over real yandex webdav")
start := time.Now()
if err := storage.Put(ctx, name, payload); err != nil {
t.Fatalf("Put: %v", err)
}
t.Logf("Put took %v", time.Since(start))
start = time.Now()
names, err := storage.List(ctx, "streams/live/c2s")
if err != nil {
t.Fatalf("List: %v", err)
}
t.Logf("List took %v", time.Since(start))
if len(names) != 1 || names[0].Name != "000000000.seg" {
t.Fatalf("List returned %v, want one segment", names)
}
start = time.Now()
got, err := storage.Get(ctx, name)
if err != nil {
t.Fatalf("Get: %v", err)
}
t.Logf("Get took %v", time.Since(start))
if !bytes.Equal(got, payload) {
t.Fatalf("Get returned %q", got)
}
if _, err := storage.Get(ctx, "streams/live/c2s/000000009.seg"); err != errNotFound {
t.Fatalf("Get of a missing object returned %v, want errNotFound", err)
}
if err := storage.Delete(ctx, name); err != nil {
t.Fatalf("Delete: %v", err)
}
}
func TestLiveYandexTransport(t *testing.T) {
settings, _, cleanup := liveYandexSettings(t)
defer cleanup()
client, server, done := pairWith(t, settings)
defer done()
start := time.Now()
if _, err := client.Write([]byte("ping")); err != nil {
t.Fatalf("client write: %v", err)
}
expectRead(t, server, "ping")
t.Logf("client to server round took %v", time.Since(start))
size := 1000000
if raw := os.Getenv("XDRIVE_LIVE_BYTES"); raw != "" {
fmt.Sscanf(raw, "%d", &size)
}
payload := make([]byte, size)
rand.Read(payload)
start = time.Now()
go func() { client.Write(payload) }()
if err := server.SetReadDeadline(time.Now().Add(5 * time.Minute)); 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)
}
elapsed := time.Since(start)
if !bytes.Equal(got, payload) {
t.Fatal("payload mismatch")
}
t.Logf("%d bytes in %v -> %.1f KiB/s", len(payload), elapsed,
float64(len(payload))/1024/elapsed.Seconds())
}