From 2157f74340a2193176669bde54ed0d65a86b4689 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Mon, 9 Oct 2023 14:14:00 +0000
Subject: [PATCH] feat: add shebang support to scripts

This enables much greater portability!
---
 agent/agentssh/agentssh.go      | 24 ++++++++++++++++++++-
 agent/agentssh/agentssh_test.go | 37 +++++++++++++++++++++++++++++++++
 dogfood/main.tf                 |  3 +++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
index 238fa55329a8c..ad45b5bd5aaca 100644
--- a/agent/agentssh/agentssh.go
+++ b/agent/agentssh/agentssh.go
@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/gliderlabs/ssh"
+	"github.com/kballard/go-shellquote"
 	"github.com/pkg/sftp"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/spf13/afero"
@@ -515,8 +516,29 @@ func (s *Server) CreateCommand(ctx context.Context, script string, env []string)
 	if runtime.GOOS == "windows" {
 		caller = "/c"
 	}
+	name := shell
 	args := []string{caller, script}
 
+	if strings.HasPrefix(strings.TrimSpace(script), "#!") {
+		// If the script starts with a shebang, we should
+		// execute it directly. This is useful for running
+		// scripts that aren't executable.
+		shebang := strings.SplitN(script, "\n", 2)[0]
+		shebang = strings.TrimPrefix(shebang, "#!")
+		shebang = strings.TrimSpace(shebang)
+		words, err := shellquote.Split(shebang)
+		if err != nil {
+			return nil, xerrors.Errorf("split shebang: %w", err)
+		}
+		name = words[0]
+		if len(words) > 1 {
+			args = words[1:]
+		} else {
+			args = []string{}
+		}
+		args = append(args, caller, script)
+	}
+
 	// gliderlabs/ssh returns a command slice of zero
 	// when a shell is requested.
 	if len(script) == 0 {
@@ -528,7 +550,7 @@ func (s *Server) CreateCommand(ctx context.Context, script string, env []string)
 		}
 	}
 
-	cmd := pty.CommandContext(ctx, shell, args...)
+	cmd := pty.CommandContext(ctx, name, args...)
 	cmd.Dir = manifest.Directory
 
 	// If the metadata directory doesn't exist, we run the command
diff --git a/agent/agentssh/agentssh_test.go b/agent/agentssh/agentssh_test.go
index 146da9b4c3bec..b72da96e4ce43 100644
--- a/agent/agentssh/agentssh_test.go
+++ b/agent/agentssh/agentssh_test.go
@@ -6,6 +6,7 @@ import (
 	"bytes"
 	"context"
 	"net"
+	"runtime"
 	"strings"
 	"sync"
 	"testing"
@@ -71,6 +72,42 @@ func TestNewServer_ServeClient(t *testing.T) {
 	<-done
 }
 
+func TestNewServer_ExecuteShebang(t *testing.T) {
+	t.Parallel()
+	if runtime.GOOS == "windows" {
+		t.Skip("bash doesn't exist on Windows")
+	}
+
+	ctx := context.Background()
+	logger := slogtest.Make(t, nil)
+	s, err := agentssh.NewServer(ctx, logger, prometheus.NewRegistry(), afero.NewMemMapFs(), 0, "")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = s.Close()
+	})
+	s.AgentToken = func() string { return "" }
+	s.Manifest = atomic.NewPointer(&agentsdk.Manifest{})
+
+	t.Run("Basic", func(t *testing.T) {
+		t.Parallel()
+		cmd, err := s.CreateCommand(ctx, `#!/bin/bash
+		echo test`, nil)
+		require.NoError(t, err)
+		output, err := cmd.AsExec().CombinedOutput()
+		require.NoError(t, err)
+		require.Equal(t, "test\n", string(output))
+	})
+	t.Run("Args", func(t *testing.T) {
+		t.Parallel()
+		cmd, err := s.CreateCommand(ctx, `#!/usr/bin/env bash
+		echo test`, nil)
+		require.NoError(t, err)
+		output, err := cmd.AsExec().CombinedOutput()
+		require.NoError(t, err)
+		require.Equal(t, "test\n", string(output))
+	})
+}
+
 func TestNewServer_CloseActiveConnections(t *testing.T) {
 	t.Parallel()
 
diff --git a/dogfood/main.tf b/dogfood/main.tf
index 94549ba11b64f..02a2346984c3a 100644
--- a/dogfood/main.tf
+++ b/dogfood/main.tf
@@ -172,6 +172,7 @@ resource "coder_agent" "dev" {
     display_name = "Swap Usage (Host)"
     key          = "4_swap_usage_host"
     script       = <<EOT
+      #!/bin/bash
       echo "$(free -b | awk '/^Swap/ { printf("%.1f/%.1f", $3/1024.0/1024.0/1024.0, $2/1024.0/1024.0/1024.0) }') GiB"
     EOT
     interval     = 10
@@ -183,6 +184,7 @@ resource "coder_agent" "dev" {
     key          = "5_load_host"
     # get load avg scaled by number of cores
     script   = <<EOT
+      #!/bin/bash
       echo "`cat /proc/loadavg | awk '{ print $1 }'` `nproc`" | awk '{ printf "%0.2f", $1/$2 }'
     EOT
     interval = 60
@@ -201,6 +203,7 @@ resource "coder_agent" "dev" {
     display_name = "Word of the Day"
     key          = "7_word"
     script       = <<EOT
+      #!/bin/bash
       curl -o - --silent https://www.merriam-webster.com/word-of-the-day 2>&1 | awk ' $0 ~ "Word of the Day: [A-z]+" { print $5; exit }'
     EOT
     interval     = 86400