sashakoshka/go-gemini
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

server: Disallow ServeConn usage after Shutdown

Browse Source
This commit is contained in:
Adnan Maolood 2021-02-24 19:22:26 -05:00
parent 2d7f28e152
commit 21ad3a2ded
1 changed files with 13 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
server.go
View File

@ -282,14 +282,17 @@ func (srv *Server) serve(ctx context.Context, l net.Listener) error {
return err return err
} }
tempDelay = 0 tempDelay = 0
go srv.ServeConn(ctx, rw) go srv.serveConn(ctx, rw, false)
} }
} }
func (srv *Server) trackConn(conn *net.Conn, cancel context.CancelFunc) bool { func (srv *Server) trackConn(conn *net.Conn, cancel context.CancelFunc, external bool) bool {
srv.mu.Lock() srv.mu.Lock()
defer srv.mu.Unlock() defer srv.mu.Unlock()
if srv.closed && !srv.shutdown { // Reject the connection under the following conditions:
// - Shutdown or Close has been called and conn is external (from ServeConn)
// - Close (not Shutdown) has been called and conn is internal (from Serve)
if srv.closed && (external || !srv.shutdown) {
return false return false
} }
if srv.conns == nil { if srv.conns == nil {
@ -309,15 +312,17 @@ func (srv *Server) deleteConn(conn *net.Conn) {
// It closes the connection when the response has been completed. // It closes the connection when the response has been completed.
// If the provided context expires before the response has completed, // If the provided context expires before the response has completed,
// ServeConn closes the connection and returns the context's error. // ServeConn closes the connection and returns the context's error.
//
// Note that ServeConn can be used during a Shutdown.
func (srv *Server) ServeConn(ctx context.Context, conn net.Conn) error { func (srv *Server) ServeConn(ctx context.Context, conn net.Conn) error {
return srv.serveConn(ctx, conn, true)
}
func (srv *Server) serveConn(ctx context.Context, conn net.Conn, external bool) error {
defer conn.Close() defer conn.Close()
ctx, cancel := context.WithCancel(ctx) ctx, cancel := context.WithCancel(ctx)
defer cancel() defer cancel()
if !srv.trackConn(&conn, cancel) { if !srv.trackConn(&conn, cancel, external) {
return context.Canceled return context.Canceled
} }
defer srv.tryCloseDone() defer srv.tryCloseDone()
@ -332,7 +337,7 @@ func (srv *Server) ServeConn(ctx context.Context, conn net.Conn) error {
errch := make(chan error, 1) errch := make(chan error, 1)
go func() { go func() {
errch <- srv.serveConn(ctx, conn) errch <- srv.goServeConn(ctx, conn)
}() }()
select { select {
@ -343,7 +348,7 @@ func (srv *Server) ServeConn(ctx context.Context, conn net.Conn) error {
} }
} }
func (srv *Server) serveConn(ctx context.Context, conn net.Conn) error { func (srv *Server) goServeConn(ctx context.Context, conn net.Conn) error {
ctx, cancel := context.WithCancel(ctx) ctx, cancel := context.WithCancel(ctx)
done := ctx.Done() done := ctx.Done()
cw := &contextWriter{ cw := &contextWriter{