1 files changed, 6 insertions, 0 deletions

diff --git a/main.go b/main.go
index ec4170b..00fe063 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+        "context"
         "errors"
         "fmt"
         "log"
@@ -34,9 +35,14 @@ var clientCmd = &cobra.Command{
         Run: func(cmd *cobra.Command, args []string) {
                 listenOn := cmd.Flag("listen").Value.String()
 
+                // TODO: Handle signals
+                ctx, cancel := context.WithCancel(context.Background())
+                defer cancel()
+
                 h := &ClientHandler{
                         SocketListenOn:  listenOn,
                         WebsocketServer: args[0],
+                        Context:         ctx,
                 }
 
                 log.Printf("Serving on %s", listenOn)