From 6d477487d93fc18c5077cdf1341ac5b44c4a0fa3 Mon Sep 17 00:00:00 2001
From: Vasile Popescu <elisescu@elisescu.com>
Date: Sun, 20 Dec 2020 19:06:06 +0100
Subject: [PATCH] Add several TTY_SHARE_* session environment variables

Set these variables in the started child shell process:
- TTY_SHARE=1 to signal running inside a tty-share session
- TTY_SHARE_LOCAL_URL - contains the URL of the public session
- TTY_SHARE_PUBLIC_URL - contains the URL of the local session
---
 main.go       | 35 ++++++++++++++++++++++++++---------
 pty_master.go |  3 ++-
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/main.go b/main.go
index 550ce62..cff1728 100644
--- a/main.go
+++ b/main.go
@@ -125,14 +125,8 @@ Flags:
 		os.Exit(1)
 	}
 
-	ptyMaster := ptyMasterNew()
-	err := ptyMaster.Start(*commandName, strings.Fields(*commandArgs))
-	if err != nil {
-		log.Errorf("Cannot start the %s command: %s", *commandName, err.Error())
-		return
-	}
-
-	sessionID := "local"
+	sessionID := ""
+	publicURL := ""
 	if *publicSession {
 		proxy, err := proxy.NewProxyConnection(*listenAddress, *proxyServerAddress, *noTLS)
 		if err != nil {
@@ -142,12 +136,35 @@ Flags:
 
 		go proxy.RunProxy()
 		sessionID = proxy.SessionID
-		fmt.Printf("public session: %s\n", proxy.PublicURL)
+		publicURL = proxy.PublicURL
 		defer proxy.Stop()
 	}
 
+	envVars := os.Environ()
+	envVars = append(envVars,
+		fmt.Sprintf("TTY_SHARE_LOCAL_URL=http://%s", *listenAddress),
+		fmt.Sprintf("TTY_SHARE=1", os.Getpid()),
+	)
+
+	if publicURL != "" {
+		envVars = append(envVars,
+			fmt.Sprintf("TTY_SHARE_PUBLIC_URL=%s", publicURL),
+		)
+	}
+
+	ptyMaster := ptyMasterNew()
+	err := ptyMaster.Start(*commandName, strings.Fields(*commandArgs), envVars)
+	if err != nil {
+		log.Errorf("Cannot start the %s command: %s", *commandName, err.Error())
+		return
+	}
+
 	// Display the session information to the user, before showing any output from the command.
 	// Wait until the user presses Enter
+	if publicURL != "" {
+		fmt.Printf("public session: %s\n", publicURL)
+	}
+
 	fmt.Printf("local session: http://%s/s/local/\n", *listenAddress)
 	fmt.Printf("Press Enter to continue!\n")
 	bufio.NewReader(os.Stdin).ReadString('\n')
diff --git a/pty_master.go b/pty_master.go
index ed0c9b2..002c389 100644
--- a/pty_master.go
+++ b/pty_master.go
@@ -30,8 +30,9 @@ func isStdinTerminal() bool {
 	return terminal.IsTerminal(0)
 }
 
-func (pty *ptyMaster) Start(command string, args []string) (err error) {
+func (pty *ptyMaster) Start(command string, args []string, envVars []string) (err error) {
 	pty.command = exec.Command(command, args...)
+	pty.command.Env = envVars
 	pty.ptyFile, err = ptyDevice.Start(pty.command)
 
 	if err != nil {