Slightly less cantankerous mini TCP echo server

tcp-challenge/server/cmd/server/main.go

@@ -1,14 +1,18 @@
 package main
 
 import (
+	"bufio"
 	"context"
-	"io"
+	"errors"
+	"fmt"
 	"log"
 	"net"
 	"os"
 	"os/signal"
 	"strconv"
+	"strings"
 	"syscall"
+	"time"
 )
 
 func main() {
@@ -36,6 +40,13 @@ func main() {
 		log.Fatal(err)
 	}
 
+	go func() {
+		<-ctx.Done()
+		if err := listener.Close(); err != nil {
+			log.Printf("failed to close listener: %v", err)
+		}
+	}()
+
 	workers := []*worker{}
 	requests := make(chan net.Conn)
 
@@ -49,15 +60,24 @@ func main() {
 		select {
 		case <-ctx.Done():
 			return
+		case <-time.After(time.Microsecond):
+			continue
 		default:
-			log.Println("accepting connection")
+			log.Printf("entering accept condition")
+
 			conn, err := listener.Accept()
 			if err != nil {
+				if errors.Is(err, net.ErrClosed) {
+					return
+				}
+
 				log.Printf("failed to accept connection: %v", err)
 				continue
 			}
 
-			log.Println("accepted connection")
+			conn.SetDeadline(time.Now().Add(5 * time.Second))
+
+			log.Printf("accepted connection %v", conn.RemoteAddr())
 			requests <- conn
 		}
 	}
@@ -77,9 +97,21 @@ func (w *worker) Start(ctx context.Context) {
 			log.Printf("[worker %v] exiting", w.id)
 			return
 		case conn := <-w.r:
-			log.Printf("[worker %v] handling connection", w.id)
+			log.Printf("[worker %v] handling connection %v", w.id, conn.RemoteAddr())
 
-			io.Copy(conn, conn)
+			fmt.Fprintf(conn, "ITS ECHO TIME\n")
+
+			scanner := bufio.NewScanner(conn)
+			for scanner.Scan() {
+				req := strings.TrimSpace(scanner.Text())
+
+				if req == "BYE" {
+					fmt.Fprintf(conn, "BYE BYE NOW\n")
+					break
+				}
+
+				fmt.Fprintf(conn, "YOU SAID: %v\n", req)
+			}
 
 			if err := conn.Close(); err != nil {
 				log.Printf("[worker %v] failed to echo: %v", w.id, err)