func connect(addr string) error { nc, err := net.Dial("tcp", *serverAddr) if err != nil { return err } defer nc.Close() c := batt.NewConn(nc) err = c.Write(batt.Message{"hello", url.Values{ "k": []string{batt.Secret}, "p": platforms, }}) if err != nil { return err } m, err := c.Read() if err != nil { return err } if m.Verb != "hello" { return fmt.Errorf(`expected "hello", got %q`, m.Verb) } errc := make(chan error, 1) done := make(chan bool, 1) go func() { for { m, err := c.Read() if err != nil { errc <- err return } select { case in <- m: case <-done: return } } }() go func() { for { var m batt.Message select { case m = <-out: case <-done: return } if err := c.Write(m); err != nil { errc <- err return } } }() err = <-errc done <- true // shut down other goroutine return err }
func handleWorkerConn(nc net.Conn) { defer nc.Close() addr := nc.RemoteAddr() boot := make(chan string, 1) defer func() { var reason string select { case reason = <-boot: default: } log.Printf("Worker conn %s disconnected; reason=%v", addr, reason) }() log.Printf("Got potential worker connection from %s", addr) // They get 5 seconds to authenticate. authTimer := time.AfterFunc(5*time.Second, func() { boot <- "login timeout" nc.Close() }) c := batt.NewConn(nc) m, err := c.Read() if err != nil { boot <- fmt.Sprintf("inital message read error: %v", err) return } if m.Verb != "hello" { boot <- "speaking while not authenticated" return } if batt.Secret == "" || m.Get("k") != batt.Secret { boot <- fmt.Sprintf("bad password; want %q", batt.Secret) return } authTimer.Stop() platforms := m.Values["p"] log.Printf("Worker conn %s authenticated; working for clients: %+v", addr, platforms) c.Write(batt.Message{Verb: "hello"}) w := &Worker{ Addr: addr.String(), Platforms: platforms, Conn: c, in: make(chan interface{}), } registerWorker(w) defer unregisterWorker(w) defer close(w.in) go w.loop() for { m, err := c.Read() if err != nil { boot <- fmt.Sprintf("message read error: %v", err) log.Printf("Worker conn %v shut down: %v", addr, err) return } w.in <- m } panic("unreachable") }