Cleanup, add local client

1 files changed, 29 insertions, 0 deletions

diff --git a/main.go b/main.go
index 00fe063..1ac3351 100644
--- a/main.go
+++ b/main.go
@@ -50,6 +50,32 @@ var clientCmd = &cobra.Command{
         },
 }
 
+var localClientCmd = &cobra.Command{
+        Use:   "localclient [server host]",
+        Short: "Act as a client for a websocket-proxy server",
+        Args: func(cmd *cobra.Command, args []string) error {
+                if len(args) != 1 || args[0] == "" {
+                        return errors.New("Server host is a required argument")
+                }
+                if !strings.HasPrefix(args[0], "ws://") && !strings.HasPrefix(args[0], "wss://") {
+                        return errors.New("Server host format is ws[s]://host[:port]/[path]")
+                }
+                return nil
+        },
+        Run: func(cmd *cobra.Command, args []string) {
+                // TODO: Handle signals
+                ctx, cancel := context.WithCancel(context.Background())
+                defer cancel()
+
+                h := &LocalClientHandler{
+                        WebsocketServer: args[0],
+                        Context:         ctx,
+                }
+
+                h.Run()
+        },
+}
+
 var serverCmd = &cobra.Command{
         Use:   "server [next-hop host]",
         Short: "Serve websocket proxy client",
@@ -69,7 +95,10 @@ var serverCmd = &cobra.Command{
 }
 
 func main() {
+        log.SetOutput(os.Stderr)
+
         rootCmd.AddCommand(clientCmd)
+        rootCmd.AddCommand(localClientCmd)
         rootCmd.AddCommand(serverCmd)
 
         clientCmd.Flags().StringP("listen", "l", ":9013", "[address]:port to bind for serving clients")