diff --git a/hysteria/service.go b/hysteria/service.go index c36ee35..5e5b69a 100644 --- a/hysteria/service.go +++ b/hysteria/service.go @@ -56,6 +56,9 @@ type Service[U comparable] struct { udpTimeout time.Duration handler ServerHandler quicListener io.Closer + + sessionAccess sync.Mutex + sessions map[*serverSession[U]]struct{} } func NewService[U comparable](options ServiceOptions) (*Service[U], error) { @@ -94,6 +97,7 @@ func NewService[U comparable](options ServiceOptions) (*Service[U], error) { handler: options.Handler, udpDisabled: options.UDPDisabled, udpTimeout: options.UDPTimeout, + sessions: make(map[*serverSession[U]]struct{}), }, nil } @@ -103,6 +107,48 @@ func (s *Service[U]) UpdateUsers(userList []U, passwordList []string) { userMap[passwordList[i]] = user } s.userMap = userMap + s.closeRemovedSessions(userList) +} + +// trackSession registers an authenticated session, so that a later UpdateUsers +// can revoke it. Only authenticated sessions are registered: before that there +// is no user to match a removed one against. +func (s *Service[U]) trackSession(session *serverSession[U]) { + s.sessionAccess.Lock() + defer s.sessionAccess.Unlock() + s.sessions[session] = struct{}{} +} + +func (s *Service[U]) untrackSession(session *serverSession[U]) { + s.sessionAccess.Lock() + defer s.sessionAccess.Unlock() + delete(s.sessions, session) +} + +// closeRemovedSessions closes the sessions of users that are gone from the new +// user table. A session authenticates once and then multiplexes every later +// stream over the same QUIC connection without consulting the user table +// again, so dropping a user from it does not by itself revoke the access that +// user already has. +// +// Victims are collected under the lock and closed after it is released, so +// that the close path is never run with the session lock held. +func (s *Service[U]) closeRemovedSessions(userList []U) { + users := make(map[U]struct{}, len(userList)) + for _, user := range userList { + users[user] = struct{}{} + } + s.sessionAccess.Lock() + var removed []*serverSession[U] + for session := range s.sessions { + if _, stillValid := users[session.authUser]; !stillValid { + removed = append(removed, session) + } + } + s.sessionAccess.Unlock() + for _, session := range removed { + session.closeWithError(E.New("user removed")) + } } func (s *Service[U]) Start(conn net.PacketConn) error { @@ -160,6 +206,7 @@ type serverSession[U comparable] struct { } func (s *serverSession[U]) handleConnection() { + defer s.untrackSession(s) ctx, cancel := context.WithTimeout(s.ctx, ProtocolTimeout) controlStream, err := s.quicConn.AcceptStream(ctx) cancel() @@ -193,6 +240,7 @@ func (s *serverSession[U]) handleConnection() { } _ = controlStream.SetDeadline(time.Time{}) s.authUser = user + s.trackSession(s) s.quicConn.SetCongestionControl(hyCC.NewBrutalSender(min(s.sendBPS, clientHello.RecvBPS), s.quicConn.InitialPacketSize(), s.brutalDebug, s.logger)) if !s.udpDisabled { go s.loopMessages() diff --git a/hysteria2/service.go b/hysteria2/service.go index e79582f..823803f 100644 --- a/hysteria2/service.go +++ b/hysteria2/service.go @@ -76,6 +76,9 @@ type Service[U comparable] struct { quicListener io.Closer bbrProfile congestion_meta2.Profile realmServer *realm.Server + + sessionAccess sync.Mutex + sessions map[*serverSession[U]]struct{} } func NewService[U comparable](options ServiceOptions) (*Service[U], error) { @@ -145,6 +148,7 @@ func NewService[U comparable](options ServiceOptions) (*Service[U], error) { masqueradeHandler: options.MasqueradeHandler, bbrProfile: bbrProfile, realmServer: realmServer, + sessions: make(map[*serverSession[U]]struct{}), }, nil } @@ -154,6 +158,48 @@ func (s *Service[U]) UpdateUsers(userList []U, passwordList []string) { userMap[passwordList[i]] = user } s.userMap = userMap + s.closeRemovedSessions(userList) +} + +// trackSession registers an authenticated session, so that a later UpdateUsers +// can revoke it. Only authenticated sessions are registered: before that there +// is no user to match a removed one against. +func (s *Service[U]) trackSession(session *serverSession[U]) { + s.sessionAccess.Lock() + defer s.sessionAccess.Unlock() + s.sessions[session] = struct{}{} +} + +func (s *Service[U]) untrackSession(session *serverSession[U]) { + s.sessionAccess.Lock() + defer s.sessionAccess.Unlock() + delete(s.sessions, session) +} + +// closeRemovedSessions closes the sessions of users that are gone from the new +// user table. A session authenticates once and then multiplexes every later +// stream over the same QUIC connection without consulting the user table +// again, so dropping a user from it does not by itself revoke the access that +// user already has. +// +// Victims are collected under the lock and closed after it is released, so +// that the close path is never run with the session lock held. +func (s *Service[U]) closeRemovedSessions(userList []U) { + users := make(map[U]struct{}, len(userList)) + for _, user := range userList { + users[user] = struct{}{} + } + s.sessionAccess.Lock() + var removed []*serverSession[U] + for session := range s.sessions { + if _, stillValid := users[session.authUser]; !stillValid { + removed = append(removed, session) + } + } + s.sessionAccess.Unlock() + for _, session := range removed { + session.closeWithError(E.New("user removed")) + } } func (s *Service[U]) Start(conn net.PacketConn) error { @@ -247,6 +293,7 @@ func (s *Service[U]) handleConnection(connection *quic.Conn) { connDone: make(chan struct{}), udpConnMap: make(map[uint32]*udpPacketConn), } + defer s.untrackSession(session) httpServer := http3.Server{ Handler: session, StreamDispatcher: session.dispatchStream, @@ -287,6 +334,7 @@ func (s *serverSession[U]) ServeHTTP(w http.ResponseWriter, r *http.Request) { } s.authUser = user s.authenticated = true + s.trackSession(s) var rxAuto bool if s.receiveBPS > 0 && s.ignoreClientBandwidth && request.Rx == 0 { s.logger.Debug("process connection from ", r.RemoteAddr, ": BBR disabled by server") diff --git a/tuic/service.go b/tuic/service.go index a4d0741..8f4ecc6 100644 --- a/tuic/service.go +++ b/tuic/service.go @@ -57,6 +57,9 @@ type Service[U comparable] struct { udpTimeout time.Duration handler ServiceHandler + sessionAccess sync.Mutex + sessions map[*serverSession[U]]struct{} + quicListener io.Closer } @@ -94,6 +97,7 @@ func NewService[U comparable](options ServiceOptions) (*Service[U], error) { authTimeout: options.AuthTimeout, udpTimeout: options.UDPTimeout, handler: options.Handler, + sessions: make(map[*serverSession[U]]struct{}), }, nil } @@ -106,6 +110,48 @@ func (s *Service[U]) UpdateUsers(userList []U, uuidList [][16]byte, passwordList } s.userMap = userMap s.passwordMap = passwordMap + s.closeRemovedSessions(userList) +} + +// trackSession registers an authenticated session, so that a later UpdateUsers +// can revoke it. Only authenticated sessions are registered: before that there +// is no user to match a removed one against. +func (s *Service[U]) trackSession(session *serverSession[U]) { + s.sessionAccess.Lock() + defer s.sessionAccess.Unlock() + s.sessions[session] = struct{}{} +} + +func (s *Service[U]) untrackSession(session *serverSession[U]) { + s.sessionAccess.Lock() + defer s.sessionAccess.Unlock() + delete(s.sessions, session) +} + +// closeRemovedSessions closes the sessions of users that are gone from the new +// user table. A session authenticates once and then multiplexes every later +// stream over the same QUIC connection without consulting the user table +// again, so dropping a user from it does not by itself revoke the access that +// user already has. +// +// Victims are collected under the lock and closed after it is released: +// closeWithError untracks the session, which takes the same lock. +func (s *Service[U]) closeRemovedSessions(userList []U) { + users := make(map[U]struct{}, len(userList)) + for _, user := range userList { + users[user] = struct{}{} + } + s.sessionAccess.Lock() + var removed []*serverSession[U] + for session := range s.sessions { + if _, stillValid := users[session.authUser]; !stillValid { + removed = append(removed, session) + } + } + s.sessionAccess.Unlock() + for _, session := range removed { + session.closeWithError(E.New("user removed")) + } } func (s *Service[U]) Start(conn net.PacketConn) error { @@ -258,6 +304,9 @@ func (s *serverSession[U]) handleUniStream(stream *quic.ReceiveStream) error { return E.New("authentication: token mismatch") } s.authUser = user + // Before authDone, so that the session cannot start relaying streams + // while it is still invisible to closeRemovedSessions. + s.trackSession(s) close(s.authDone) return nil case CommandPacket: @@ -390,6 +439,7 @@ func (s *serverSession[U]) closeWithError(err error) { s.connErr = err close(s.connDone) } + s.untrackSession(s) if E.IsClosedOrCanceled(err) { s.logger.Debug(E.Cause(err, "connection failed")) } else { diff --git a/tuic/service_test.go b/tuic/service_test.go new file mode 100644 index 0000000..f803bfc --- /dev/null +++ b/tuic/service_test.go @@ -0,0 +1,191 @@ +package tuic + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "math/big" + "net" + "testing" + "time" + + "github.com/sagernet/sing/common/logger" + M "github.com/sagernet/sing/common/metadata" + N "github.com/sagernet/sing/common/network" + aTLS "github.com/sagernet/sing/common/tls" + + "github.com/gofrs/uuid/v5" +) + +// A session authenticates once and then relays every later stream without +// consulting the user table again, so removing a user from it has to close +// that user's sessions as well. +func TestUpdateUsersClosesRemovedUserSessions(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + userUUID := uuid.Must(uuid.NewV4()) + handler := &testHandler{connections: make(chan net.Conn, 1)} + certificate := generateCertificate(t) + + service, err := NewService[string](ServiceOptions{ + Context: ctx, + Logger: logger.NOP(), + TLSConfig: &testTLSConfig{config: &tls.Config{ + Certificates: []tls.Certificate{certificate}, + NextProtos: []string{"tuic-test"}, + }}, + Handler: handler, + }) + if err != nil { + t.Fatal(err) + } + service.UpdateUsers([]string{"alice"}, [][16]byte{userUUID}, []string{"password"}) + + packetConn, err := net.ListenPacket("udp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + err = service.Start(packetConn) + if err != nil { + t.Fatal(err) + } + defer service.Close() + + client, err := NewClient(ClientOptions{ + Context: ctx, + Dialer: N.SystemDialer, + ServerAddress: M.SocksaddrFromNet(packetConn.LocalAddr()), + TLSConfig: &testTLSConfig{config: &tls.Config{ + ServerName: "example.org", + InsecureSkipVerify: true, + NextProtos: []string{"tuic-test"}, + }}, + UUID: userUUID, + Password: "password", + }) + if err != nil { + t.Fatal(err) + } + + clientConn, err := client.DialConn(ctx, M.ParseSocksaddr("example.org:80")) + if err != nil { + t.Fatal(err) + } + // The request header rides on the first write, so nothing reaches the + // handler until the client sends something. + _, err = clientConn.Write([]byte("ping")) + if err != nil { + t.Fatal(err) + } + select { + case <-handler.connections: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for the relayed connection") + } + + service.UpdateUsers(nil, nil, nil) + + readDone := make(chan error, 1) + go func() { + _, readErr := clientConn.Read(make([]byte, 1)) + readDone <- readErr + }() + select { + case readErr := <-readDone: + if readErr == nil { + t.Fatal("expected the relayed connection to fail") + } + case <-time.After(10 * time.Second): + t.Fatal("session of a removed user stayed alive") + } +} + +type testHandler struct { + connections chan net.Conn +} + +func (h *testHandler) NewConnectionEx(ctx context.Context, conn net.Conn, source M.Socksaddr, destination M.Socksaddr, onClose N.CloseHandlerFunc) { + h.connections <- conn +} + +func (h *testHandler) NewPacketConnectionEx(ctx context.Context, conn N.PacketConn, source M.Socksaddr, destination M.Socksaddr, onClose N.CloseHandlerFunc) { + conn.Close() +} + +func generateCertificate(t *testing.T) tls.Certificate { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatal(err) + } + template := x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "example.org"}, + DNSNames: []string{"example.org"}, + NotBefore: time.Now().Add(-time.Hour), + NotAfter: time.Now().Add(time.Hour), + } + certificate, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key) + if err != nil { + t.Fatal(err) + } + return tls.Certificate{Certificate: [][]byte{certificate}, PrivateKey: key} +} + +type testTLSConfig struct { + config *tls.Config +} + +func (c *testTLSConfig) ServerName() string { + return c.config.ServerName +} + +func (c *testTLSConfig) SetServerName(serverName string) { + c.config.ServerName = serverName +} + +func (c *testTLSConfig) NextProtos() []string { + return c.config.NextProtos +} + +func (c *testTLSConfig) SetNextProtos(nextProto []string) { + c.config.NextProtos = nextProto +} + +func (c *testTLSConfig) HandshakeTimeout() time.Duration { + return 0 +} + +func (c *testTLSConfig) SetHandshakeTimeout(timeout time.Duration) { +} + +func (c *testTLSConfig) STDConfig() (*aTLS.STDConfig, error) { + return c.config, nil +} + +func (c *testTLSConfig) Client(conn net.Conn) (aTLS.Conn, error) { + return tls.Client(conn, c.config), nil +} + +func (c *testTLSConfig) Server(conn net.Conn) (aTLS.Conn, error) { + return tls.Server(conn, c.config), nil +} + +func (c *testTLSConfig) Clone() aTLS.Config { + return &testTLSConfig{config: c.config.Clone()} +} + +func (c *testTLSConfig) Start() error { + return nil +} + +func (c *testTLSConfig) Close() error { + return nil +}