add websocket

api/node_type.go

@@ -15,3 +15,27 @@ type Node struct {
  Type string `json:"type"`
  UpTime int `json:"uptime"`
 }
+
+type TermProxy struct {
+ Port string `json:"port"`
+ Ticket string `json:"ticket"`
+ UPID string `json:"upid"`
+ User string `json:"user"`
+}
+
+type TermProxyOption struct {
+ CMD string `json:"cmd,omitempty"`
+ CMDOpts string `json:"cmd-opts,omitempty"`
+}
+
+type VNCShellOption struct {
+ TermProxyOption
+ Height int `json:"height,omitempty"`
+ Websocket bool `json:"websocket,omitempty"`
+ Width int `json:"width,omitempty"`
+}
+
+type VNCWebSocket struct {
+ Port string `json:"port,omitempty"`
+ VNCTicket string `json:"vncticket,omitempty"`
+}

go.mod

@@ -3,6 +3,7 @@ module github.com/sp-yduck/proxmox-go
 go 1.20
 
 require (
+ github.com/gorilla/websocket v1.5.0
  github.com/pkg/errors v0.9.1
  github.com/stretchr/testify v1.8.4
 )

go.sum

@@ -1,5 +1,7 @@
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
+github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

proxmox/websocket.go

@@ -0,0 +1,120 @@
+package proxmox
+
+import (
+ "context"
+ "fmt"
+ "regexp"
+ "strconv"
+ "strings"
+
+ "github.com/pkg/errors"
+
+ "github.com/gorilla/websocket"
+ "github.com/sp-yduck/proxmox-go/api"
+)
+
+const (
+ finMessage = "done with status: "
+ finMessageFormat = finMessage + `[0-9]+`
+)
+
+type VNCWebSocketClient struct {
+ conn *websocket.Conn
+}
+
+func (s *Service) NewNodeVNCWebSocketConnection(ctx context.Context, nodeName string) (*VNCWebSocketClient, error) {
+ termProxy, err := s.restclient.CreateNodeTermProxy(ctx, nodeName, api.TermProxyOption{})
+ if err != nil {
+ return nil, err
+ }
+ conn, err := s.restclient.DialNodeVNCWebSocket(ctx, nodeName, *termProxy)
+ if err != nil {
+ return nil, err
+ }
+
+ return &VNCWebSocketClient{conn: conn}, nil
+}
+
+func (c *VNCWebSocketClient) Close() {
+ c.conn.Close()
+}
+
+func (c *VNCWebSocketClient) Write(cmd string) error {
+ b := []byte(fmt.Sprintf("%s\n", cmd))
+ bheader := []byte(fmt.Sprintf("0:%d:", len(b)))
+ bmsg := append(bheader, b...)
+ if err := c.conn.WriteMessage(websocket.BinaryMessage, bmsg); err != nil {
+ return err
+ }
+ return c.sendFinMessage()
+}
+
+func (c *VNCWebSocketClient) sendFinMessage() error {
+ b := []byte(fmt.Sprintf(`echo "%s$?"%s`, finMessage, "\n"))
+ bheader := []byte(fmt.Sprintf("0:%d:", len(b)))
+ bmsg := append(bheader, b...)
+ if err := c.conn.WriteMessage(websocket.BinaryMessage, bmsg); err != nil {
+ return err
+ }
+ return nil
+}
+
+// Read() reads message until find fin message
+// then returns whole message and status code
+func (c *VNCWebSocketClient) Read(ctx context.Context) (outputs string, code int, err error) {
+ done := make(chan error, 1)
+ go func() {
+ defer close(done)
+ for {
+ _, msg, err := c.conn.ReadMessage()
+ if err != nil {
+ done <- err
+ return
+ }
+ outputs += string(msg)
+ finMsg := parseFinMessage(string(msg))
+ if finMsg != "" {
+ code, err = parseStatusFromFinMessage(finMsg)
+ done <- err
+ return
+ }
+ }
+ }()
+ select {
+ case err = <-done:
+ return outputs, code, err
+ case <-ctx.Done():
+ return outputs, -1, errors.New("context deadline exceeded")
+ }
+}
+
+// Exec executes a command and return error if code is not 0
+// usually out contains many extra messages that is just useless
+func (c *VNCWebSocketClient) Exec(ctx context.Context, cmd string) (out string, code int, err error) {
+ if err := c.Write(cmd); err != nil {
+ return "", 0, err
+ }
+ out, code, err = c.Read(ctx)
+ if err != nil {
+ return out, code, err
+ }
+ if code != 0 {
+ return out, code, errors.Errorf("exit with non zero code: %d", code)
+ }
+ return out, 0, nil
+}
+
+func parseFinMessage(message string) string {
+ re := regexp.MustCompile(finMessageFormat)
+ return re.FindString(message)
+}
+
+func parseStatusFromFinMessage(message string) (int, error) {
+ re := regexp.MustCompile(finMessageFormat)
+ match := re.FindString(message)
+ if match == "" {
+ return 0, errors.Errorf("failed to find status code from %s", message)
+ }
+ statusCode := strings.Split(match, ": ")[1]
+ return strconv.Atoi(statusCode)
+}

proxmox/websocket_test.go

@@ -0,0 +1,51 @@
+package proxmox
+
+import (
+ "context"
+ "testing"
+ "time"
+)
+
+func (s *TestSuite) TestVNCWebSocketClient() {
+ testNode := s.getTestNode()
+ client, err := s.service.NewNodeVNCWebSocketConnection(context.TODO(), testNode.Node)
+ if err != nil {
+ s.T().Fatalf("failed to create new vnc client: %v", err)
+ }
+ defer client.Close()
+
+ if err := client.Write("pwd"); err != nil {
+ s.T().Fatalf("write error: %v", err)
+ }
+
+ ctx, _ := context.WithTimeout(context.TODO(), 10*time.Second)
+ out, _, err := client.Read(ctx)
+ if err != nil {
+ s.T().Fatalf("failed read message: %v", err)
+ }
+
+ s.T().Logf("read message: %s", out)
+}
+
+func (s *TestSuite) TestExec() {
+ testNode := s.getTestNode()
+ client, err := s.service.NewNodeVNCWebSocketConnection(context.TODO(), testNode.Node)
+ if err != nil {
+ s.T().Fatalf("failed to create new vnc client: %v", err)
+ }
+ defer client.Close()
+
+ ctx, _ := context.WithTimeout(context.TODO(), 5*time.Second)
+ out, code, err := client.Exec(ctx, "whoami | base64 | base64 -d")
+ if err != nil {
+ s.T().Fatalf("failed to exec command: %s : %d : %v", out, code, err)
+ }
+ s.T().Logf("exec command : %s : %d", out, code)
+}
+
+func TestParseFinMessage(t *testing.T) {
+ testMsg := " daf" + finMessage + "123\n"
+ if parseFinMessage(testMsg) == "" {
+ t.Fatalf("wrong")
+ }
+}

rest/client.go

@@ -8,7 +8,6 @@ import (
  "io"
  "io/ioutil"
  "net/http"
- "net/url"
  "strings"
  "time"
 
@@ -62,6 +61,7 @@ func complementURL(url string) string {
  if !strings.HasPrefix(url, "http") {
  url = "http://" + url
  }
+ url, _ = strings.CutSuffix(url, "/")
  return url
 }
 
@@ -96,10 +96,7 @@ func WithAPIToken(tokenid, secret string) ClientOption {
 }
 
 func (c *RESTClient) Do(ctx context.Context, httpMethod, urlPath string, req, v interface{}) error {
- url, err := url.JoinPath(c.endpoint, urlPath)
- if err != nil {
- return err
- }
+ endpoint := c.endpoint + urlPath
 
  var body io.Reader
  if req != nil {
@@ -110,7 +107,7 @@ func (c *RESTClient) Do(ctx context.Context, httpMethod, urlPath string, req, v
  body = bytes.NewReader(jsonReq)
  }
 
- httpReq, err := http.NewRequestWithContext(ctx, httpMethod, url, body)
+ httpReq, err := http.NewRequestWithContext(ctx, httpMethod, endpoint, body)
  if err != nil {
  return err
  }
@@ -162,7 +159,6 @@ func (c *RESTClient) Delete(ctx context.Context, path string, req, res interface
 
 func (c *RESTClient) makeAuthHeaders() http.Header {
  header := make(http.Header)
- // header.Add("User-Agent", c.userAgent)
  header.Add("Accept", "application/json")
  if c.token != "" {
  header.Add("Authorization", fmt.Sprintf("PVEAPIToken=%s", c.token))

rest/node.go

@@ -2,6 +2,8 @@ package rest
 
 import (
  "context"
+ "fmt"
+ "net/url"
 
  "github.com/sp-yduck/proxmox-go/api"
 )
@@ -26,3 +28,30 @@ func (c *RESTClient) GetNode(ctx context.Context, name string) (*api.Node, error
  }
  return nil, NotFoundErr
 }
+
+func (c *RESTClient) CreateNodeTermProxy(ctx context.Context, nodeName string, option api.TermProxyOption) (*api.TermProxy, error) {
+ path := fmt.Sprintf("/nodes/%s/termproxy", nodeName)
+ var termProxy *api.TermProxy
+ if err := c.Post(ctx, path, option, &termProxy); err != nil {
+ return nil, err
+ }
+ return termProxy, nil
+}
+
+func (c *RESTClient) CreateNodeVNCShell(ctx context.Context, nodeName string, option api.VNCShellOption) (*api.TermProxy, error) {
+ path := fmt.Sprintf("/nodes/%s/vncshell", nodeName)
+ var termProxy *api.TermProxy
+ if err := c.Post(ctx, path, option, &termProxy); err != nil {
+ return nil, err
+ }
+ return termProxy, nil
+}
+
+func (c *RESTClient) GetNodeVNCWebSocket(ctx context.Context, nodeName, port, vncticket string) (*api.VNCWebSocket, error) {
+ path := fmt.Sprintf("/nodes/%s/vncwebsocket?port=%s&vncticket=%s", nodeName, port, url.QueryEscape(vncticket))
+ var websocket *api.VNCWebSocket
+ if err := c.Get(ctx, path, &websocket); err != nil {
+ return nil, err
+ }
+ return websocket, nil
+}

rest/nodes_test.go

@@ -0,0 +1,31 @@
+package rest
+
+import (
+ "context"
+
+ "github.com/sp-yduck/proxmox-go/api"
+)
+
+func (s *TestSuite) TestCreateTermProxy() {
+ testNode := s.GetTestNode()
+ termProxy, err := s.restclient.CreateNodeTermProxy(context.TODO(), testNode.Node, api.TermProxyOption{})
+ if err != nil {
+ s.T().Fatalf("failed to create termproxy: %v", err)
+ }
+ s.T().Logf("create termproxy: %v", termProxy)
+}
+
+func (s *TestSuite) TestGetVNCWebSocket() {
+ testNode := s.GetTestNode()
+ termProxy, err := s.restclient.CreateNodeTermProxy(context.TODO(), testNode.Node, api.TermProxyOption{})
+ if err != nil {
+ s.T().Fatalf("failed to create termproxy: %v", err)
+ }
+ s.T().Logf("create termproxy: %v", termProxy)
+
+ websocket, err := s.restclient.GetNodeVNCWebSocket(context.TODO(), testNode.Node, termProxy.Port, termProxy.Ticket)
+ if err != nil {
+ s.T().Fatalf("failed to get vncwebsocket: %v", err)
+ }
+ s.T().Logf("get vncwebsocket: %v", websocket)
+}

rest/websocket.go

@@ -0,0 +1,49 @@
+package rest
+
+import (
+ "context"
+ "crypto/tls"
+ "fmt"
+ "net/http"
+ "net/url"
+ "strings"
+ "time"
+
+ "github.com/gorilla/websocket"
+ "github.com/pkg/errors"
+
+ "github.com/sp-yduck/proxmox-go/api"
+)
+
+func (c *RESTClient) DialNodeVNCWebSocket(ctx context.Context, nodeName string, vnc api.TermProxy) (*websocket.Conn, error) {
+ baseUrl := strings.Replace(c.endpoint, "https://", "wss://", 1)
+ baseUrl = strings.Replace(baseUrl, "http://", "wss://", 1)
+ websocketUrl := fmt.Sprintf("%s/nodes/%s/vncwebsocket?port=%s&vncticket=%s", baseUrl, nodeName, vnc.Port, url.QueryEscape(vnc.Ticket))
+
+ conn, resp, err := c.websocketDialer().DialContext(ctx, websocketUrl, c.makeAuthHeaders())
+ if err != nil {
+ if resp != nil {
+ return nil, errors.Errorf("failed to dial websocket: %v : %v", checkResponse(resp), err)
+ }
+ return nil, errors.Errorf("failed to dial websocket: %v", err)
+ }
+
+ if err := conn.WriteMessage(websocket.BinaryMessage, []byte(fmt.Sprintf("%s:%s\n", vnc.User, vnc.Ticket))); err != nil {
+ return nil, errors.Errorf("failed to start session: %v", err)
+ }
+
+ return conn, nil
+}
+
+func (c *RESTClient) websocketDialer() *websocket.Dialer {
+ var tlsConfig *tls.Config
+ transport := c.httpClient.Transport.(*http.Transport)
+ if transport != nil {
+ tlsConfig = transport.TLSClientConfig
+ }
+ return &websocket.Dialer{
+ Proxy: http.ProxyFromEnvironment,
+ HandshakeTimeout: 30 * time.Second,
+ TLSClientConfig: tlsConfig,
+ }
+}