Buffer fixing and simplification

[ucspi.git] / cmd / tlss / main.go

diff --git a/cmd/tlss/main.go b/cmd/tlss/main.go
index 7740ecd5cdf07443908de03885378f48437120e1..881d2210990b85d1d745b5b6ac0f0c78e749c2fd 100644 (file)
--- a/cmd/tlss/main.go
+++ b/cmd/tlss/main.go
@@ -22,6 +22,7 @@ import (
        "crypto/x509"
        "flag"
        "fmt"
+       "io"
        "log"
        "os"
        "os/exec"
@@ -70,7 +71,7 @@ func main() {
                cfg.ClientAuth = tls.RequireAndVerifyClientCert
        }
 
-       conn := &ucspi.Conn{R: os.Stdin, W: os.Stdout}
+       conn, _ := ucspi.NewConn(os.Stdin, os.Stdout)
        tlsConn := tls.Server(conn, cfg)
        if err = tlsConn.Handshake(); err != nil {
                log.Fatalln(err)
@@ -80,10 +81,18 @@ func main() {
                dn = tlsConn.ConnectionState().PeerCertificates[0].Subject.String()
        }
 
+       rr, rw, err := os.Pipe()
+       if err != nil {
+               log.Fatalln(err)
+       }
+       wr, ww, err := os.Pipe()
+       if err != nil {
+               log.Fatalln(err)
+       }
        args := flag.Args()
        cmd := exec.Command(args[0], args[1:]...)
-       cmd.Stdin = tlsConn
-       cmd.Stdout = tlsConn
+       cmd.Stdin = rr
+       cmd.Stdout = ww
        cmd.Stderr = os.Stderr
        cmd.Env = append(os.Environ(), "PROTO=TLS")
        if dn != "" {
@@ -93,7 +102,19 @@ func main() {
        if err = cmd.Start(); err != nil {
                log.Fatalln(err)
        }
-       if _, err = cmd.Process.Wait(); err != nil {
+       worker := make(chan struct{})
+       go func() {
+               io.Copy(rw, tlsConn)
+       }()
+       go func() {
+               io.Copy(tlsConn, wr)
+               tlsConn.Close()
+               close(worker)
+       }()
+       err = cmd.Wait()
+       ww.Close()
+       <-worker
+       if err != nil {
                log.Fatalln(err)
        }
 }