Single port graceful shutdown

connection.go

@@ -40,7 +40,6 @@ type chanConnection struct {
 	channel       chan []byte
 	srcAddr, addr *net.UDPAddr
 	timeout       time.Duration
-	complete      chan string
 }
 
 func (c *chanConnection) sendTo(data []byte, addr *net.UDPAddr) error {
@@ -67,9 +66,7 @@ func (c *chanConnection) sendTo(data []byte, addr *net.UDPAddr) error {
 func (c *chanConnection) readFrom(buffer []byte) (int, *net.UDPAddr, error) {
 	select {
 	case data := <-c.channel:
-		for i := range data {
-			buffer[i] = data[i]
-		}
+		copy(buffer, data)
 		return len(data), c.addr, nil
 	case <-time.After(c.timeout):
 		return 0, nil, makeError(c.addr.String())
@@ -98,7 +95,7 @@ func makeError(addr string) net.Error {
 		timeout:   true,
 		temporary: true,
 	}
-	error.error = fmt.Errorf("Channel timeout: %v", addr)
+	error.error = fmt.Errorf("channel timeout: %v", addr)
 	return &error
 }
 

netascii/netascii_test.go

@@ -2,7 +2,7 @@ package netascii
 
 import (
 	"bytes"
-	"io/ioutil"
+	"io"
 	"strings"
 	"testing"
 	"testing/iotest"
@@ -20,7 +20,7 @@ var basic = map[string]string{
 func TestTo(t *testing.T) {
 	for text, netascii := range basic {
 		to := ToReader(strings.NewReader(text))
-		n, _ := ioutil.ReadAll(to)
+		n, _ := io.ReadAll(to)
 		if !bytes.Equal(n, []byte(netascii)) {
 			t.Errorf("%q to netascii: %q != %q", text, n, netascii)
 		}
@@ -33,7 +33,7 @@ func TestFrom(t *testing.T) {
 		b := &bytes.Buffer{}
 		from := FromWriter(b)
 		r.WriteTo(from)
-		n, _ := ioutil.ReadAll(b)
+		n, _ := io.ReadAll(b)
 		if string(n) != text {
 			t.Errorf("%q from netascii: %q != %q", netascii, n, text)
 		}
@@ -69,7 +69,7 @@ func TestWriteRead(t *testing.T) {
 	two := &bytes.Buffer{}
 	from := FromWriter(two)
 	one.WriteTo(from)
-	text2, _ := ioutil.ReadAll(two)
+	text2, _ := io.ReadAll(two)
 	if text != string(text2) {
 		t.Errorf("text mismatch \n%x \n%x", text, text2)
 	}
@@ -82,7 +82,7 @@ func TestOneByte(t *testing.T) {
 	two := &bytes.Buffer{}
 	from := FromWriter(two)
 	one.WriteTo(from)
-	text2, _ := ioutil.ReadAll(two)
+	text2, _ := io.ReadAll(two)
 	if text != string(text2) {
 		t.Errorf("text mismatch \n%x \n%x", text, text2)
 	}

server.go

@@ -208,7 +208,7 @@ func (s *Server) Serve(conn net.PacketConn) error {
 	// but necessary at this point.
 	addr := net.ParseIP(host)
 	if addr == nil {
-		return fmt.Errorf("Failed to determine IP class of listening address")
+		return fmt.Errorf("failed to determine IP class of listening address")
 	}
 
 	if conn, ok := s.conn.(*net.UDPConn); ok {

single_port.go

@@ -2,46 +2,66 @@ package tftp
 
 import (
 	"net"
+	"time"
 )
 
 func (s *Server) singlePortProcessRequests() error {
+	shuttingDown := false
 	for {
 		select {
 		case <-s.cancel.Done():
-			s.wg.Wait()
-			return nil
+			shuttingDown = true
 		default:
-			buf := make([]byte, s.maxBlockLen+4)
-			cnt, localAddr, srcAddr, maxSz, err := s.getPacket(buf)
-			if err != nil || cnt == 0 {
-				if s.hook != nil {
+		}
+
+		if shuttingDown {
+			// So we not blocked forever waiting for a packet
+			s.conn.SetReadDeadline(time.Now().Add(time.Second))
+		}
+
+		buf := make([]byte, s.maxBlockLen+4)
+		cnt, localAddr, srcAddr, maxSz, err := s.getPacket(buf)
+		if err != nil {
+			if shuttingDown {
+				s.wg.Wait()
+				return nil
+			}
+			if s.hook != nil {
+				s.hook.OnFailure(TransferStats{
+					SenderAnticipateEnabled: s.sendAEnable,
+				}, err)
+			}
+			continue
+		}
+		if cnt == 0 {
+			continue
+		}
+		s.Lock()
+		if receiverChannel, ok := s.handlers[srcAddr.String()]; ok {
+			// Packet received for a transfer in progress.
+			s.Unlock()
+			select {
+			case receiverChannel <- buf[:cnt]:
+			default:
+				// We don't want to block the main loop if a channel is full
+			}
+		} else {
+			// No existing transfer for given source address. Start a new one.
+			if shuttingDown {
+				s.Unlock()
+				continue
+			}
+			lc := make(chan []byte, 1)
+			s.handlers[srcAddr.String()] = lc
+			s.Unlock()
+			go func() {
+				err := s.handlePacket(localAddr, srcAddr, buf, cnt, maxSz, lc)
+				if err != nil && s.hook != nil {
 					s.hook.OnFailure(TransferStats{
 						SenderAnticipateEnabled: s.sendAEnable,
 					}, err)
 				}
-				continue
-			}
-			s.Lock()
-			if receiverChannel, ok := s.handlers[srcAddr.String()]; ok {
-				s.Unlock()
-				select {
-				case receiverChannel <- buf[:cnt]:
-				default:
-					// We don't want to block the main loop if a channel is full
-				}
-			} else {
-				lc := make(chan []byte, 1)
-				s.handlers[srcAddr.String()] = lc
-				s.Unlock()
-				go func() {
-					err := s.handlePacket(localAddr, srcAddr, buf, cnt, maxSz, lc)
-					if err != nil && s.hook != nil {
-						s.hook.OnFailure(TransferStats{
-							SenderAnticipateEnabled: s.sendAEnable,
-						}, err)
-					}
-				}()
-			}
+			}()
 		}
 	}
 }