# HG changeset patch
# User Henry Precheur <henry@precheur.org>
# Date 1489808707 25200
#      Fri Mar 17 20:45:07 2017 -0700
# Node ID 943347af9ed359a92e578926ce4a7131223185fd
# Parent  e919bba7f4666876f0a1f1179f1d36f05c03ffbd
fix keepalive

diff --git a/cmd/kgp_client/main.go b/cmd/kgp_client/main.go
--- a/cmd/kgp_client/main.go
+++ b/cmd/kgp_client/main.go
@@ -4,12 +4,13 @@
 	"flag"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"log"
 	"net"
 	"os"
 	"strings"
 
+	"context"
+
 	"bitbucket.org/henry/kgp"
 )
 
@@ -50,10 +51,19 @@
 	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 @@
 	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()))
 }
diff --git a/kgp.go b/kgp.go
--- a/kgp.go
+++ b/kgp.go
@@ -12,7 +12,7 @@
 	"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 @@
 	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 @@
 	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 @@
 		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 @@
 				}
 			}
 		case p := <-input:
+			// Connection was closed
+			if p == nil {
+				break
+			}
 			switch p.Type {
 			case keepalivePacket:
 				go keepAlive(done, p, output)