Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Server able to close all connections for given user #125

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Server able to close all connections for given user #125

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 32 additions & 10 deletions server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ type Server struct {
listenerWg sync.WaitGroup
mu sync.Mutex
listeners map[net.Listener]struct{}
conns map[*gossh.ServerConn]struct{}
userConns map[string]map[*gossh.ServerConn]struct{}
connWg sync.WaitGroup
doneChan chan struct{}
}
Expand Down Expand Up @@ -155,9 +155,12 @@ func (srv *Server) Close() error {
defer srv.mu.Unlock()
srv.closeDoneChanLocked()
err := srv.closeListenersLocked()
for c := range srv.conns {
c.Close()
delete(srv.conns, c)
for user, conns := range srv.userConns {
for c := range conns {
c.Close()
delete(conns, c)
}
delete(srv.userConns, user)
}
return err
}
Expand Down Expand Up @@ -323,6 +326,19 @@ func (srv *Server) SetOption(option Option) error {
return option(srv)
}

func (srv *Server) CloseConnectionsForUser(user string) {
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is too specific to your use-case. A better way to generalize this would be to make connections exported or a hook to get all connections or something, so the specific implementations can do whatever they want with all the connections.

srv.mu.Lock()
defer srv.mu.Unlock()

if srv.userConns[user] == nil {
return
}

for conn := range srv.userConns[user] {
conn.Close()
}
}

func (srv *Server) getDoneChan() <-chan struct{} {
srv.mu.Lock()
defer srv.mu.Unlock()
Expand Down Expand Up @@ -368,7 +384,7 @@ func (srv *Server) trackListener(ln net.Listener, add bool) {
if add {
// If the *Server is being reused after a previous
// Close or Shutdown, reset its doneChan:
if len(srv.listeners) == 0 && len(srv.conns) == 0 {
if len(srv.listeners) == 0 && len(srv.userConns) == 0 {
srv.doneChan = nil
}
srv.listeners[ln] = struct{}{}
Expand All @@ -382,14 +398,20 @@ func (srv *Server) trackListener(ln net.Listener, add bool) {
func (srv *Server) trackConn(c *gossh.ServerConn, add bool) {
srv.mu.Lock()
defer srv.mu.Unlock()
if srv.conns == nil {
srv.conns = make(map[*gossh.ServerConn]struct{})
if srv.userConns == nil {
srv.userConns = make(map[string]map[*gossh.ServerConn]struct{})
}
if srv.userConns[c.User()] == nil {
srv.userConns[c.User()] = make(map[*gossh.ServerConn]struct{})
}
if add {
srv.conns[c] = struct{}{}
srv.userConns[c.User()][c] = struct{}{}
srv.connWg.Add(1)
} else {
delete(srv.conns, c)
delete(srv.userConns[c.User()], c)
if len(srv.userConns[c.User()]) == 0 {
delete(srv.userConns, c.User())
}
srv.connWg.Done()
}
}
}