From feacb1f85e29769c2e0ff18e3705ef2bfc55556c Mon Sep 17 00:00:00 2001 From: Toby Date: Sat, 15 Jun 2024 14:15:56 -0700 Subject: [PATCH] feat(core): server RequestHook support --- app/go.mod | 4 - app/go.sum | 4 - core/go.sum | 10 ++ core/internal/integration_tests/.mockery.yaml | 3 + core/internal/integration_tests/hook_test.go | 144 ++++++++++++++++++ .../mocks/mock_RequestHook.go | 141 +++++++++++++++++ core/server/config.go | 13 ++ core/server/mock_udpIO.go | 47 ++++++ core/server/server.go | 24 ++- core/server/udp.go | 30 +++- core/server/udp_test.go | 1 + extras/go.mod | 5 +- extras/go.sum | 9 ++ 13 files changed, 416 insertions(+), 19 deletions(-) create mode 100644 core/internal/integration_tests/hook_test.go create mode 100644 core/internal/integration_tests/mocks/mock_RequestHook.go diff --git a/app/go.mod b/app/go.mod index 46cf92c..bfc20b0 100644 --- a/app/go.mod +++ b/app/go.mod @@ -33,7 +33,6 @@ require ( github.com/fsnotify/fsnotify v1.7.0 // indirect github.com/go-ole/go-ole v1.3.0 // indirect github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect - github.com/golang/protobuf v1.5.4 // indirect github.com/google/go-querystring v1.1.0 // indirect github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 // indirect github.com/hashicorp/go-cleanhttp v0.5.2 // indirect @@ -42,7 +41,6 @@ require ( github.com/hashicorp/hcl v1.0.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/klauspost/cpuid/v2 v2.1.1 // indirect - github.com/kr/pretty v0.3.1 // indirect github.com/libdns/libdns v0.2.2 // indirect github.com/magiconair/properties v1.8.7 // indirect github.com/miekg/dns v1.1.59 // indirect @@ -53,7 +51,6 @@ require ( github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/quic-go/qpack v0.4.0 // indirect - github.com/rogpeppe/go-internal v1.12.0 // indirect github.com/sagernet/netlink v0.0.0-20220905062125-8043b4a9aa97 // indirect github.com/scjalliance/comshim v0.0.0-20230315213746-5e51f40bd3b9 // indirect github.com/spf13/afero v1.9.3 // indirect @@ -77,7 +74,6 @@ require ( golang.org/x/text v0.15.0 // indirect golang.org/x/tools v0.21.0 // indirect google.golang.org/protobuf v1.34.1 // indirect - gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect gopkg.in/ini.v1 v1.67.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect rsc.io/qr v0.2.0 // indirect diff --git a/app/go.sum b/app/go.sum index 8f07498..550394f 100644 --- a/app/go.sum +++ b/app/go.sum @@ -59,7 +59,6 @@ github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGX github.com/cncf/udpa/go v0.0.0-20200629203442-efcf912fb354/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk= github.com/cncf/udpa/go v0.0.0-20201120205902-5459f2c99403/go.mod h1:WmhPx2Nbnhtbo57+VJT5O0JRkEi1Wbu0z5j0R8u5Hbk= github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= -github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -169,7 +168,6 @@ github.com/klauspost/cpuid/v2 v2.1.1 h1:t0wUqjowdm8ezddV5k0tLWVklVuvLJpoHeb4WBdy github.com/klauspost/cpuid/v2 v2.1.1/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY= github.com/kr/fs v0.1.0/go.mod h1:FFnZGqtBN9Gxj7eW1uZ42v5BccTP0vu6NEaFoC2HwRg= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= -github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= @@ -215,7 +213,6 @@ github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaR github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ= github.com/pelletier/go-toml/v2 v2.0.6 h1:nrzqCb7j9cDFj2coyLNLaZuJTLjWjlaz6nvTvIwycIU= github.com/pelletier/go-toml/v2 v2.0.6/go.mod h1:eumQOmlWiOPt5WriQQqoM5y18pDHwha2N+QD+EUNTek= -github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= @@ -226,7 +223,6 @@ github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1: github.com/quic-go/qpack v0.4.0 h1:Cr9BXA1sQS2SmDUWjSofMPNKmvF6IiIfDRmgU0w1ZCo= github.com/quic-go/qpack v0.4.0/go.mod h1:UZVnYIfi5GRk+zI9UMaCPsmZ2xKJP7XBUvVyT1Knj9A= github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4= -github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= diff --git a/core/go.sum b/core/go.sum index 96fd988..a8e10c4 100644 --- a/core/go.sum +++ b/core/go.sum @@ -12,12 +12,15 @@ github.com/go-logr/logr v1.2.4/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbV github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI= github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 h1:yAJXTCF9TqKcTiHJAE8dj7HMvPfh66eeA2JYW7eFpSE= github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE= github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= @@ -26,15 +29,20 @@ github.com/onsi/ginkgo/v2 v2.9.5 h1:+6Hr4uxzP4XIUyAkg61dWBw8lb/gc4/X5luuxN/EC+Q= github.com/onsi/ginkgo/v2 v2.9.5/go.mod h1:tvAoo1QUJwNEU2ITftXTpR7R1RbCzoZUOs3RonqW57k= github.com/onsi/gomega v1.27.6 h1:ENqfyGeS5AX/rlXDd/ETokDz93u0YufY1Pgxuy/PvWE= github.com/onsi/gomega v1.27.6/go.mod h1:PIQNjfQwkP3aQAH7lf7j87O/5FiNr+ZR8+ipb+qQlhg= +github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/quic-go/qpack v0.4.0 h1:Cr9BXA1sQS2SmDUWjSofMPNKmvF6IiIfDRmgU0w1ZCo= github.com/quic-go/qpack v0.4.0/go.mod h1:UZVnYIfi5GRk+zI9UMaCPsmZ2xKJP7XBUvVyT1Knj9A= +github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= +github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= go.uber.org/goleak v1.2.1 h1:NBol2c7O1ZokfZ0LEU9K6Whx/KnwvepVetCUhtKja4A= go.uber.org/goleak v1.2.1/go.mod h1:qlT2yGI9QafXHhZZLxlSuNsMw3FFLxBr+tBRlmO1xH4= go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU= @@ -59,8 +67,10 @@ golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= golang.org/x/tools v0.21.0 h1:qc0xYgIbsSDt9EyWz05J5wfa7LOVW0YTLOXrqdLAWIw= golang.org/x/tools v0.21.0/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk= google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg= +google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/core/internal/integration_tests/.mockery.yaml b/core/internal/integration_tests/.mockery.yaml index ad233ba..550a725 100644 --- a/core/internal/integration_tests/.mockery.yaml +++ b/core/internal/integration_tests/.mockery.yaml @@ -24,3 +24,6 @@ packages: TrafficLogger: config: mockname: MockTrafficLogger + RequestHook: + config: + mockname: MockRequestHook \ No newline at end of file diff --git a/core/internal/integration_tests/hook_test.go b/core/internal/integration_tests/hook_test.go new file mode 100644 index 0000000..14f67bc --- /dev/null +++ b/core/internal/integration_tests/hook_test.go @@ -0,0 +1,144 @@ +package integration_tests + +import ( + "io" + "net" + "testing" + + "github.com/apernet/hysteria/core/v2/client" + "github.com/apernet/hysteria/core/v2/internal/integration_tests/mocks" + "github.com/apernet/hysteria/core/v2/server" + "github.com/apernet/quic-go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" +) + +func TestClientServerHookTCP(t *testing.T) { + fakeEchoAddr := "hahanope:6666" + realEchoAddr := "127.0.0.1:22333" + + // Create server + udpConn, udpAddr, err := serverConn() + assert.NoError(t, err) + auth := mocks.NewMockAuthenticator(t) + auth.EXPECT().Authenticate(mock.Anything, mock.Anything, mock.Anything).Return(true, "nobody") + hook := mocks.NewMockRequestHook(t) + hook.EXPECT().TCP(mock.Anything, mock.Anything).RunAndReturn(func(stream quic.Stream, s *string) ([]byte, error) { + assert.Equal(t, fakeEchoAddr, *s) + // Change the address + *s = realEchoAddr + // Read the first 5 bytes and replace them with "byeee" + data := make([]byte, 5) + _, err := io.ReadFull(stream, data) + if err != nil { + return nil, err + } + assert.Equal(t, []byte("hello"), data) + return []byte("byeee"), nil + }).Once() + s, err := server.NewServer(&server.Config{ + TLSConfig: serverTLSConfig(), + Conn: udpConn, + RequestHook: hook, + Authenticator: auth, + }) + assert.NoError(t, err) + defer s.Close() + go s.Serve() + + // Create TCP echo server + echoListener, err := net.Listen("tcp", realEchoAddr) + assert.NoError(t, err) + echoServer := &tcpEchoServer{Listener: echoListener} + defer echoServer.Close() + go echoServer.Serve() + + // Create client + c, _, err := client.NewClient(&client.Config{ + ServerAddr: udpAddr, + TLSConfig: client.TLSConfig{InsecureSkipVerify: true}, + FastOpen: true, // Client MUST have FastOpen for this + }) + assert.NoError(t, err) + defer c.Close() + + // Dial TCP + conn, err := c.TCP(fakeEchoAddr) + assert.NoError(t, err) + defer conn.Close() + + // Send and receive data + sData := []byte("hello world") + _, err = conn.Write(sData) + assert.NoError(t, err) + rData := make([]byte, len(sData)) + _, err = io.ReadFull(conn, rData) + assert.NoError(t, err) + assert.Equal(t, []byte("byeee world"), rData) +} + +func TestClientServerHookUDP(t *testing.T) { + fakeEchoAddr := "hahanope:6666" + realEchoAddr := "127.0.0.1:22333" + + // Create server + udpConn, udpAddr, err := serverConn() + assert.NoError(t, err) + auth := mocks.NewMockAuthenticator(t) + auth.EXPECT().Authenticate(mock.Anything, mock.Anything, mock.Anything).Return(true, "nobody") + hook := mocks.NewMockRequestHook(t) + hook.EXPECT().UDP(mock.Anything, mock.Anything).RunAndReturn(func(bytes []byte, s *string) error { + assert.Equal(t, fakeEchoAddr, *s) + assert.Equal(t, []byte("hello world"), bytes) + // Change the address + *s = realEchoAddr + return nil + }).Once() + s, err := server.NewServer(&server.Config{ + TLSConfig: serverTLSConfig(), + Conn: udpConn, + RequestHook: hook, + Authenticator: auth, + }) + assert.NoError(t, err) + defer s.Close() + go s.Serve() + + // Create UDP echo server + echoConn, err := net.ListenPacket("udp", realEchoAddr) + assert.NoError(t, err) + echoServer := &udpEchoServer{Conn: echoConn} + defer echoServer.Close() + go echoServer.Serve() + + // Create client + c, _, err := client.NewClient(&client.Config{ + ServerAddr: udpAddr, + TLSConfig: client.TLSConfig{InsecureSkipVerify: true}, + }) + assert.NoError(t, err) + defer c.Close() + + // Listen UDP + conn, err := c.UDP() + assert.NoError(t, err) + defer conn.Close() + + // Send and receive data + sData := []byte("hello world") + err = conn.Send(sData, fakeEchoAddr) + assert.NoError(t, err) + rData, rAddr, err := conn.Receive() + assert.NoError(t, err) + assert.Equal(t, sData, rData) + assert.Equal(t, realEchoAddr, rAddr) + + // Subsequent packets should also be sent to the real echo server + sData = []byte("never stop fighting") + err = conn.Send(sData, fakeEchoAddr) + assert.NoError(t, err) + rData, rAddr, err = conn.Receive() + assert.NoError(t, err) + assert.Equal(t, sData, rData) + assert.Equal(t, realEchoAddr, rAddr) +} diff --git a/core/internal/integration_tests/mocks/mock_RequestHook.go b/core/internal/integration_tests/mocks/mock_RequestHook.go new file mode 100644 index 0000000..a7a6cc5 --- /dev/null +++ b/core/internal/integration_tests/mocks/mock_RequestHook.go @@ -0,0 +1,141 @@ +// Code generated by mockery v2.43.0. DO NOT EDIT. + +package mocks + +import ( + quic "github.com/apernet/quic-go" + mock "github.com/stretchr/testify/mock" +) + +// MockRequestHook is an autogenerated mock type for the RequestHook type +type MockRequestHook struct { + mock.Mock +} + +type MockRequestHook_Expecter struct { + mock *mock.Mock +} + +func (_m *MockRequestHook) EXPECT() *MockRequestHook_Expecter { + return &MockRequestHook_Expecter{mock: &_m.Mock} +} + +// TCP provides a mock function with given fields: stream, reqAddr +func (_m *MockRequestHook) TCP(stream quic.Stream, reqAddr *string) ([]byte, error) { + ret := _m.Called(stream, reqAddr) + + if len(ret) == 0 { + panic("no return value specified for TCP") + } + + var r0 []byte + var r1 error + if rf, ok := ret.Get(0).(func(quic.Stream, *string) ([]byte, error)); ok { + return rf(stream, reqAddr) + } + if rf, ok := ret.Get(0).(func(quic.Stream, *string) []byte); ok { + r0 = rf(stream, reqAddr) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]byte) + } + } + + if rf, ok := ret.Get(1).(func(quic.Stream, *string) error); ok { + r1 = rf(stream, reqAddr) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockRequestHook_TCP_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'TCP' +type MockRequestHook_TCP_Call struct { + *mock.Call +} + +// TCP is a helper method to define mock.On call +// - stream quic.Stream +// - reqAddr *string +func (_e *MockRequestHook_Expecter) TCP(stream interface{}, reqAddr interface{}) *MockRequestHook_TCP_Call { + return &MockRequestHook_TCP_Call{Call: _e.mock.On("TCP", stream, reqAddr)} +} + +func (_c *MockRequestHook_TCP_Call) Run(run func(stream quic.Stream, reqAddr *string)) *MockRequestHook_TCP_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(quic.Stream), args[1].(*string)) + }) + return _c +} + +func (_c *MockRequestHook_TCP_Call) Return(_a0 []byte, _a1 error) *MockRequestHook_TCP_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockRequestHook_TCP_Call) RunAndReturn(run func(quic.Stream, *string) ([]byte, error)) *MockRequestHook_TCP_Call { + _c.Call.Return(run) + return _c +} + +// UDP provides a mock function with given fields: data, reqAddr +func (_m *MockRequestHook) UDP(data []byte, reqAddr *string) error { + ret := _m.Called(data, reqAddr) + + if len(ret) == 0 { + panic("no return value specified for UDP") + } + + var r0 error + if rf, ok := ret.Get(0).(func([]byte, *string) error); ok { + r0 = rf(data, reqAddr) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockRequestHook_UDP_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UDP' +type MockRequestHook_UDP_Call struct { + *mock.Call +} + +// UDP is a helper method to define mock.On call +// - data []byte +// - reqAddr *string +func (_e *MockRequestHook_Expecter) UDP(data interface{}, reqAddr interface{}) *MockRequestHook_UDP_Call { + return &MockRequestHook_UDP_Call{Call: _e.mock.On("UDP", data, reqAddr)} +} + +func (_c *MockRequestHook_UDP_Call) Run(run func(data []byte, reqAddr *string)) *MockRequestHook_UDP_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].([]byte), args[1].(*string)) + }) + return _c +} + +func (_c *MockRequestHook_UDP_Call) Return(_a0 error) *MockRequestHook_UDP_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockRequestHook_UDP_Call) RunAndReturn(run func([]byte, *string) error) *MockRequestHook_UDP_Call { + _c.Call.Return(run) + return _c +} + +// NewMockRequestHook creates a new instance of MockRequestHook. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewMockRequestHook(t interface { + mock.TestingT + Cleanup(func()) +}) *MockRequestHook { + mock := &MockRequestHook{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/core/server/config.go b/core/server/config.go index c6bf733..bd8b7ca 100644 --- a/core/server/config.go +++ b/core/server/config.go @@ -8,6 +8,7 @@ import ( "github.com/apernet/hysteria/core/v2/errors" "github.com/apernet/hysteria/core/v2/internal/pmtud" + "github.com/apernet/quic-go" ) const ( @@ -22,6 +23,7 @@ type Config struct { TLSConfig TLSConfig QUICConfig QUICConfig Conn net.PacketConn + RequestHook RequestHook Outbound Outbound BandwidthConfig BandwidthConfig IgnoreClientBandwidth bool @@ -110,6 +112,17 @@ type QUICConfig struct { DisablePathMTUDiscovery bool // The server may still override this to true on unsupported platforms. } +// RequestHook allows filtering and modifying requests before the server connects to the remote. +// The returned byte slice, if not empty, will be sent to the remote before proxying - this is +// mainly for "putting back" the content read from the client for sniffing, etc. +// Return a non-nil error to abort the connection. +// Note that due to the current architectural limitations, it can only inspect the first packet +// of a UDP connection. It also cannot put back any data as the first packet is always sent as-is. +type RequestHook interface { + TCP(stream quic.Stream, reqAddr *string) ([]byte, error) + UDP(data []byte, reqAddr *string) error +} + // Outbound provides the implementation of how the server should connect to remote servers. // Although UDP includes a reqAddr, the implementation does not necessarily have to use it // to make a "connected" UDP connection that does not accept packets from other addresses. diff --git a/core/server/mock_udpIO.go b/core/server/mock_udpIO.go index 3950cb9..bbdcf94 100644 --- a/core/server/mock_udpIO.go +++ b/core/server/mock_udpIO.go @@ -20,6 +20,53 @@ func (_m *mockUDPIO) EXPECT() *mockUDPIO_Expecter { return &mockUDPIO_Expecter{mock: &_m.Mock} } +// Hook provides a mock function with given fields: data, reqAddr +func (_m *mockUDPIO) Hook(data []byte, reqAddr *string) error { + ret := _m.Called(data, reqAddr) + + if len(ret) == 0 { + panic("no return value specified for Hook") + } + + var r0 error + if rf, ok := ret.Get(0).(func([]byte, *string) error); ok { + r0 = rf(data, reqAddr) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// mockUDPIO_Hook_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Hook' +type mockUDPIO_Hook_Call struct { + *mock.Call +} + +// Hook is a helper method to define mock.On call +// - data []byte +// - reqAddr *string +func (_e *mockUDPIO_Expecter) Hook(data interface{}, reqAddr interface{}) *mockUDPIO_Hook_Call { + return &mockUDPIO_Hook_Call{Call: _e.mock.On("Hook", data, reqAddr)} +} + +func (_c *mockUDPIO_Hook_Call) Run(run func(data []byte, reqAddr *string)) *mockUDPIO_Hook_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].([]byte), args[1].(*string)) + }) + return _c +} + +func (_c *mockUDPIO_Hook_Call) Return(_a0 error) *mockUDPIO_Hook_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *mockUDPIO_Hook_Call) RunAndReturn(run func([]byte, *string) error) *mockUDPIO_Hook_Call { + _c.Call.Return(run) + return _c +} + // ReceiveMessage provides a mock function with given fields: func (_m *mockUDPIO) ReceiveMessage() (*protocol.UDPMessage, error) { ret := _m.Called() diff --git a/core/server/server.go b/core/server/server.go index bd61db0..7ebb387 100644 --- a/core/server/server.go +++ b/core/server/server.go @@ -170,7 +170,7 @@ func (h *h3sHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { if !h.config.DisableUDP { go func() { sm := newUDPSessionManager( - &udpIOImpl{h.conn, id, h.config.TrafficLogger, h.config.Outbound}, + &udpIOImpl{h.conn, id, h.config.TrafficLogger, h.config.RequestHook, h.config.Outbound}, &udpEventLoggerImpl{h.conn, id, h.config.EventLogger}, h.config.UDPIdleTimeout) h.udpSM = sm @@ -211,6 +211,16 @@ func (h *h3sHandler) handleTCPRequest(stream quic.Stream) { _ = stream.Close() return } + // Call the hook if set + var putback []byte + if h.config.RequestHook != nil { + putback, err = h.config.RequestHook.TCP(stream, &reqAddr) + if err != nil { + _ = protocol.WriteTCPResponse(stream, false, err.Error()) + _ = stream.Close() + return + } + } // Log the event if h.config.EventLogger != nil { h.config.EventLogger.TCPRequest(h.conn.RemoteAddr(), h.authID, reqAddr) @@ -227,6 +237,10 @@ func (h *h3sHandler) handleTCPRequest(stream quic.Stream) { return } _ = protocol.WriteTCPResponse(stream, true, "") + // Put back the data if the hook requested + if len(putback) > 0 { + _, _ = tConn.Write(putback) + } // Start proxying if h.config.TrafficLogger != nil { err = copyTwoWayWithLogger(h.authID, stream, tConn, h.config.TrafficLogger) @@ -260,6 +274,7 @@ type udpIOImpl struct { Conn quic.Connection AuthID string TrafficLogger TrafficLogger + RequestHook RequestHook Outbound Outbound } @@ -304,6 +319,13 @@ func (io *udpIOImpl) SendMessage(buf []byte, msg *protocol.UDPMessage) error { return io.Conn.SendDatagram(buf[:msgN]) } +func (io *udpIOImpl) Hook(data []byte, reqAddr *string) error { + if io.RequestHook != nil { + return io.RequestHook.UDP(data, reqAddr) + } + return nil +} + func (io *udpIOImpl) UDP(reqAddr string) (UDPConn, error) { return io.Outbound.UDP(reqAddr) } diff --git a/core/server/udp.go b/core/server/udp.go index 320dc00..ecaee29 100644 --- a/core/server/udp.go +++ b/core/server/udp.go @@ -20,6 +20,7 @@ const ( type udpIO interface { ReceiveMessage() (*protocol.UDPMessage, error) SendMessage([]byte, *protocol.UDPMessage) error + Hook(data []byte, reqAddr *string) error UDP(reqAddr string) (UDPConn, error) } @@ -29,11 +30,12 @@ type udpEventLogger interface { } type udpSessionEntry struct { - ID uint32 - Conn UDPConn - D *frag.Defragger - Last *utils.AtomicTime - Timeout bool // true if the session is closed due to timeout + ID uint32 + Conn UDPConn + OverrideAddr string // Ignore the address in the UDP message, always use this if not empty + D *frag.Defragger + Last *utils.AtomicTime + Timeout bool // true if the session is closed due to timeout } // Feed feeds a UDP message to the session. @@ -47,7 +49,11 @@ func (e *udpSessionEntry) Feed(msg *protocol.UDPMessage) (int, error) { if dfMsg == nil { return 0, nil } - return e.Conn.WriteTo(dfMsg.Data, dfMsg.Addr) + if e.OverrideAddr != "" { + return e.Conn.WriteTo(dfMsg.Data, e.OverrideAddr) + } else { + return e.Conn.WriteTo(dfMsg.Data, dfMsg.Addr) + } } // ReceiveLoop receives incoming UDP packets, packs them into UDP messages, @@ -177,7 +183,15 @@ func (m *udpSessionManager) feed(msg *protocol.UDPMessage) { // Create a new session if not exists if entry == nil { + // Call the hook + origMsgAddr := msg.Addr + err := m.io.Hook(msg.Data, &msg.Addr) + if err != nil { + return + } + // Log the event m.eventLogger.New(msg.SessionID, msg.Addr) + // Dial target & create a new session entry conn, err := m.io.UDP(msg.Addr) if err != nil { m.eventLogger.Close(msg.SessionID, err) @@ -189,6 +203,10 @@ func (m *udpSessionManager) feed(msg *protocol.UDPMessage) { D: &frag.Defragger{}, Last: utils.NewAtomicTime(time.Now()), } + if origMsgAddr != msg.Addr { + // Hook changed the address, enable address override + entry.OverrideAddr = msg.Addr + } // Start the receive loop for this session go func() { err := entry.ReceiveLoop(m.io) diff --git a/core/server/udp_test.go b/core/server/udp_test.go index a0169e3..aa78327 100644 --- a/core/server/udp_test.go +++ b/core/server/udp_test.go @@ -25,6 +25,7 @@ func TestUDPSessionManager(t *testing.T) { } return m, nil }) + io.EXPECT().Hook(mock.Anything, mock.Anything).Return(nil, nil) go sm.Run() diff --git a/extras/go.mod b/extras/go.mod index efc1111..5830418 100644 --- a/extras/go.mod +++ b/extras/go.mod @@ -18,14 +18,12 @@ require ( github.com/apernet/quic-go v0.44.1-0.20240520215222-bb2e53664023 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect - github.com/golang/protobuf v1.5.4 // indirect github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 // indirect - github.com/kr/pretty v0.3.1 // indirect + github.com/kr/text v0.2.0 // indirect github.com/onsi/ginkgo/v2 v2.9.5 // indirect github.com/patrickmn/go-cache v2.1.0+incompatible // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/quic-go/qpack v0.4.0 // indirect - github.com/rogpeppe/go-internal v1.12.0 // indirect github.com/stretchr/objx v0.5.2 // indirect github.com/txthinking/runnergroup v0.0.0-20210608031112-152c7c4432bf // indirect go.uber.org/mock v0.4.0 // indirect @@ -35,7 +33,6 @@ require ( golang.org/x/sys v0.20.0 // indirect golang.org/x/text v0.15.0 // indirect golang.org/x/tools v0.21.0 // indirect - gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/extras/go.sum b/extras/go.sum index 4ef8205..5d931d7 100644 --- a/extras/go.sum +++ b/extras/go.sum @@ -5,6 +5,7 @@ github.com/babolivier/go-doh-client v0.0.0-20201028162107-a76cff4cb8b6/go.mod h1 github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI= github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= @@ -13,6 +14,7 @@ github.com/go-logr/logr v1.2.4/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbV github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI= github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 h1:yAJXTCF9TqKcTiHJAE8dj7HMvPfh66eeA2JYW7eFpSE= @@ -21,10 +23,12 @@ github.com/hashicorp/golang-lru/v2 v2.0.5 h1:wW7h1TG88eUIJ2i69gaE3uNVtEPIagzhGvH github.com/hashicorp/golang-lru/v2 v2.0.5/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/miekg/dns v1.1.51/go.mod h1:2Z9d3CP1LQWihRZUf29mQ19yDThaI4DAYzte2CaQW5c= github.com/miekg/dns v1.1.59 h1:C9EXc/UToRwKLhK5wKU/I4QVsBUc8kE6MkHBkeypWZs= +github.com/miekg/dns v1.1.59/go.mod h1:nZpewl5p6IvctfgrckopVx2OlSEHPRO/U4SYkRklrEk= github.com/onsi/ginkgo/v2 v2.9.5 h1:+6Hr4uxzP4XIUyAkg61dWBw8lb/gc4/X5luuxN/EC+Q= github.com/onsi/ginkgo/v2 v2.9.5/go.mod h1:tvAoo1QUJwNEU2ITftXTpR7R1RbCzoZUOs3RonqW57k= github.com/onsi/gomega v1.27.6 h1:ENqfyGeS5AX/rlXDd/ETokDz93u0YufY1Pgxuy/PvWE= @@ -36,10 +40,13 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN github.com/quic-go/qpack v0.4.0 h1:Cr9BXA1sQS2SmDUWjSofMPNKmvF6IiIfDRmgU0w1ZCo= github.com/quic-go/qpack v0.4.0/go.mod h1:UZVnYIfi5GRk+zI9UMaCPsmZ2xKJP7XBUvVyT1Knj9A= github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= +github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/txthinking/runnergroup v0.0.0-20210608031112-152c7c4432bf h1:7PflaKRtU4np/epFxRXlFhlzLXZzKFrH5/I4so5Ove0= github.com/txthinking/runnergroup v0.0.0-20210608031112-152c7c4432bf/go.mod h1:CLUSJbazqETbaR+i0YAhXBICV9TrKH93pziccMhmhpM= github.com/txthinking/socks5 v0.0.0-20230325130024-4230056ae301 h1:d/Wr/Vl/wiJHc3AHYbYs5I3PucJvRuw3SvbmlIRf+oM= @@ -98,8 +105,10 @@ golang.org/x/tools v0.21.0 h1:qc0xYgIbsSDt9EyWz05J5wfa7LOVW0YTLOXrqdLAWIw= golang.org/x/tools v0.21.0/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg= +google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=