@@ 4,12 4,13 @@ import (
 	"flag"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"log"
 	"net"
 	"os"
 	"strings"
 
+	"context"
+
 	"bitbucket.org/henry/kgp"
 )
 

          
@@ 50,10 51,19 @@ func ParseAddress(s string) (a Address, 
 	return
 }
 
-func handshake(conn net.Conn) error {
-	var a, err = kgp.ReadAnnouncement(conn)
+func handshake(conn net.Conn, meta io.Reader) (a kgp.Announcement, err error) {
+	var my = kgp.Announcement{
+		Metadata: meta,
+	}
+	var n int
+	if n, err = my.Write(conn); err != nil {
+		fmt.Println("wrote", n)
+		return
+	}
+
+	a, err = kgp.ReadAnnouncement(conn)
 	log.Print(a)
-	return err
+	return
 }
 
 const usage = `usage: kgp_client upstream downstream

          
@@ 100,15 110,24 @@ func main() {
 	if err != nil {
 		log.Fatalln("error connecting to ", upstream, ": ", err.Error())
 	}
+	log.Print("connected")
 
-	if err = handshake(upconn); err != nil {
+	if _, err = handshake(upconn, meta); err != nil {
 		log.Fatalln(err)
 	}
-	var my = kgp.Announcement{
-		Metadata: meta,
-	}
-	if _, err := my.Write(upconn); err != nil {
-		log.Fatalln(err)
-	}
-	ioutil.ReadAll(upconn)
+
+	// var buf = make([]byte, 1000)
+	// var n int
+	// n, err = upconn.Read(buf)
+	// fmt.Println(n, err, buf)
+
+	var relay = kgp.NewRelay(
+		upconn,
+		func(c context.Context) (io.ReadWriteCloser, error) {
+			return (&net.Dialer{}).DialContext(
+				c, downstream.Network, downstream.Address,
+			)
+		},
+	)
+	fmt.Println(relay.Run(context.Background()))
 }

          

@@ 12,7 12,7 @@ import (
 	"log"
 	"sync/atomic"
 
-	"golang.org/x/net/context"
+	"context"
 )
 
 // That's from the spec, no idea why it's the limit

          
@@ 52,10 52,11 @@ func (a *Announcement) Write(w io.Writer
 	if err = binary.Write(w, binary.BigEndian, &packet); err != nil {
 		return
 	}
+	n = binary.Size(&packet)
 
 	var x int64
 	x, err = io.Copy(w, &meta)
-	n = int(x)
+	n += int(x)
 
 	return
 }

          
@@ 131,7 132,7 @@ type packetHeader struct {
 	Ack      uint32
 	Control  int16
 
-	PayloadLen uint32
+	PayloadLen uint16
 }
 
 func (packet *Packet) Write(w io.Writer, seq, ack uint32) (n int, err error) {

          
@@ 144,7 145,7 @@ func (packet *Packet) Write(w io.Writer,
 		Sequence:     seq,
 		Ack:          ack,
 		Control:      control,
-		PayloadLen:   uint32(packet.Payload.Len()),
+		PayloadLen:   uint16(packet.Payload.Len()),
 	}
 	err = binary.Write(w, binary.BigEndian, &header)
 	if err != nil {

          
@@ 471,6 472,10 @@ func (relay *Relay) Run(ctx context.Cont
 				}
 			}
 		case p := <-input:
+			// Connection was closed
+			if p == nil {
+				break
+			}
 			switch p.Type {
 			case keepalivePacket:
 				go keepAlive(done, p, output)