diff --git a/examples/echo/main.go b/examples/echo/main.go index 7297e6b10c99c52dc817ad2c24bb9c7947a3913f..cc37453da3f8416e73d525a1e3a49788385c5542 100644 --- a/examples/echo/main.go +++ b/examples/echo/main.go @@ -134,8 +134,12 @@ func main() { // a user-defined protocol name. ha.SetStreamHandler("/echo/1.0.0", func(s net.Stream) { log.Println("Got a new stream!") - defer s.Close() - doEcho(s) + if err := doEcho(s); err != nil { + log.Println(err) + s.Reset() + } else { + s.Close() + } }) if *target == "" { @@ -194,18 +198,14 @@ func main() { } // doEcho reads a line of data a stream and writes it back -func doEcho(s net.Stream) { +func doEcho(s net.Stream) error { buf := bufio.NewReader(s) str, err := buf.ReadString('\n') if err != nil { - log.Println(err) - return + return err } log.Printf("read: %s\n", str) _, err = s.Write([]byte(str)) - if err != nil { - log.Println(err) - return - } + return err }