fix socket close race

irc/socket.go

@@ -7,8 +7,15 @@ import (
 	"strings"
 )
 
+const (
+	R = '→'
+	W = '←'
+)
+
 type Socket struct {
+	closed  bool
 	conn    net.Conn
+	done    chan bool
 	reader  *bufio.Reader
 	receive chan string
 	send    chan string
@@ -18,6 +25,7 @@ type Socket struct {
 func NewSocket(conn net.Conn) *Socket {
 	socket := &Socket{
 		conn:    conn,
+		done:    make(chan bool),
 		reader:  bufio.NewReader(conn),
 		receive: make(chan string),
 		send:    make(chan string),
@@ -35,14 +43,20 @@ func (socket *Socket) String() string {
 }
 
 func (socket *Socket) Close() {
-	socket.conn.Close()
+	if socket.closed {
+		return
+	}
+
+	socket.closed = true
+	socket.done <- true
+	close(socket.done)
 }
 
 func (socket *Socket) Read() <-chan string {
 	return socket.receive
 }
 
-func (socket *Socket) Write(lines []string) {
+func (socket *Socket) Write(lines ...string) {
 	for _, line := range lines {
 		socket.send <- line
 	}
@@ -52,10 +66,7 @@ func (socket *Socket) Write(lines []string) {
 func (socket *Socket) readLines() {
 	for {
 		line, err := socket.reader.ReadString('\n')
-		if err != nil {
-			if DEBUG_NET {
-				log.Printf("%s → error: %s", socket, err)
-			}
+		if socket.isError(err, R) {
 			break
 		}
 
@@ -69,31 +80,51 @@ func (socket *Socket) readLines() {
 
 		socket.receive <- line
 	}
+
 	close(socket.receive)
-}
-
-func (socket *Socket) writeLines() {
-	for line := range socket.send {
-		if DEBUG_NET {
-			log.Printf("%s ← %s", socket, line)
-		}
-		if _, err := socket.writer.WriteString(line); socket.maybeLogWriteError(err) {
-			break
-		}
-		if _, err := socket.writer.WriteString(CRLF); socket.maybeLogWriteError(err) {
-			break
-		}
-
-		if err := socket.writer.Flush(); socket.maybeLogWriteError(err) {
-			break
-		}
+	if DEBUG_NET {
+		log.Printf("%s closed", socket)
 	}
 }
 
-func (socket *Socket) maybeLogWriteError(err error) bool {
+func (socket *Socket) writeLines() {
+	done := false
+	for !done {
+		select {
+		case line := <-socket.send:
+			if _, err := socket.writer.WriteString(line); socket.isError(err, W) {
+				break
+			}
+			if _, err := socket.writer.WriteString(CRLF); socket.isError(err, W) {
+				break
+			}
+
+			if err := socket.writer.Flush(); socket.isError(err, W) {
+				break
+			}
+			if DEBUG_NET {
+				log.Printf("%s ← %s", socket, line)
+			}
+
+		case done = <-socket.done:
+			continue
+		}
+	}
+
+	if DEBUG_NET {
+		log.Printf("%s closing", socket)
+	}
+	socket.conn.Close()
+
+	for _ = range socket.send {
+		// discard lines
+	}
+}
+
+func (socket *Socket) isError(err error, dir rune) bool {
 	if err != nil {
 		if DEBUG_NET {
-			log.Printf("%s ← error: %s", socket, err)
+			log.Printf("%s %c error: %s", socket, dir, err)
 		}
 		return true
 	}