mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
123 lines
3.1 KiB
Go
123 lines
3.1 KiB
Go
package agentssh_test
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"net"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/gliderlabs/ssh"
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/spf13/afero"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
gossh "golang.org/x/crypto/ssh"
|
|
|
|
"cdr.dev/slog"
|
|
"cdr.dev/slog/sloggers/slogtest"
|
|
"github.com/coder/coder/v2/agent/agentssh"
|
|
"github.com/coder/coder/v2/testutil"
|
|
)
|
|
|
|
func TestServer_X11(t *testing.T) {
|
|
t.Parallel()
|
|
if runtime.GOOS != "linux" {
|
|
t.Skip("X11 forwarding is only supported on Linux")
|
|
}
|
|
|
|
ctx := context.Background()
|
|
logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
|
|
fs := afero.NewOsFs()
|
|
s, err := agentssh.NewServer(ctx, logger, prometheus.NewRegistry(), fs, &agentssh.Config{})
|
|
require.NoError(t, err)
|
|
defer s.Close()
|
|
|
|
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
|
require.NoError(t, err)
|
|
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
err := s.Serve(ln)
|
|
assert.Error(t, err) // Server is closed.
|
|
}()
|
|
|
|
c := sshClient(t, ln.Addr().String())
|
|
|
|
sess, err := c.NewSession()
|
|
require.NoError(t, err)
|
|
|
|
wantScreenNumber := 1
|
|
reply, err := sess.SendRequest("x11-req", true, gossh.Marshal(ssh.X11{
|
|
AuthProtocol: "MIT-MAGIC-COOKIE-1",
|
|
AuthCookie: hex.EncodeToString([]byte("cookie")),
|
|
ScreenNumber: uint32(wantScreenNumber),
|
|
}))
|
|
require.NoError(t, err)
|
|
assert.True(t, reply)
|
|
|
|
// Want: ~DISPLAY=localhost:10.1
|
|
out, err := sess.Output("echo DISPLAY=$DISPLAY")
|
|
require.NoError(t, err)
|
|
|
|
sc := bufio.NewScanner(bytes.NewReader(out))
|
|
displayNumber := -1
|
|
for sc.Scan() {
|
|
line := strings.TrimSpace(sc.Text())
|
|
t.Log(line)
|
|
if strings.HasPrefix(line, "DISPLAY=") {
|
|
parts := strings.SplitN(line, "=", 2)
|
|
display := parts[1]
|
|
parts = strings.SplitN(display, ":", 2)
|
|
parts = strings.SplitN(parts[1], ".", 2)
|
|
displayNumber, err = strconv.Atoi(parts[0])
|
|
require.NoError(t, err)
|
|
assert.GreaterOrEqual(t, displayNumber, 10, "display number should be >= 10")
|
|
gotScreenNumber, err := strconv.Atoi(parts[1])
|
|
require.NoError(t, err)
|
|
assert.Equal(t, wantScreenNumber, gotScreenNumber, "screen number should match")
|
|
break
|
|
}
|
|
}
|
|
require.NoError(t, sc.Err())
|
|
require.NotEqual(t, -1, displayNumber)
|
|
|
|
x11Chans := c.HandleChannelOpen("x11")
|
|
payload := "hello world"
|
|
require.Eventually(t, func() bool {
|
|
conn, err := net.Dial("tcp", fmt.Sprintf("localhost:%d", agentssh.X11StartPort+displayNumber))
|
|
if err == nil {
|
|
_, err = conn.Write([]byte(payload))
|
|
assert.NoError(t, err)
|
|
_ = conn.Close()
|
|
}
|
|
return err == nil
|
|
}, testutil.WaitShort, testutil.IntervalFast)
|
|
|
|
x11 := <-x11Chans
|
|
ch, reqs, err := x11.Accept()
|
|
require.NoError(t, err)
|
|
go gossh.DiscardRequests(reqs)
|
|
got := make([]byte, len(payload))
|
|
_, err = ch.Read(got)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, payload, string(got))
|
|
_ = ch.Close()
|
|
_ = s.Close()
|
|
<-done
|
|
|
|
// Ensure the Xauthority file was written!
|
|
home, err := os.UserHomeDir()
|
|
require.NoError(t, err)
|
|
_, err = fs.Stat(filepath.Join(home, ".Xauthority"))
|
|
require.NoError(t, err)
|
|
}
|