Additional LongPollWait supporting code (#11)

Author: johnmaguire

client.go

@@ -31,12 +31,11 @@ type Client struct {
 func NewClient(useragent string, dnServer string) *Client {
 	return &Client{
 		client: &http.Client{
-			Timeout: 1 * time.Minute,
+			Timeout: 2 * time.Minute,
 			Transport: &uaTransport{
 				T: &http.Transport{
-					Proxy:                 http.ProxyFromEnvironment,
-					TLSHandshakeTimeout:   10 * time.Second,
-					ResponseHeaderTimeout: 10 * time.Second,
+					Proxy:               http.ProxyFromEnvironment,
+					TLSHandshakeTimeout: 10 * time.Second,
 					DialContext: (&net.Dialer{
 						Timeout: 10 * time.Second,
 					}).DialContext,
@@ -48,9 +47,8 @@ func NewClient(useragent string, dnServer string) *Client {
 			Timeout: 15 * time.Minute,
 			Transport: &uaTransport{
 				T: &http.Transport{
-					Proxy:                 http.ProxyFromEnvironment,
-					TLSHandshakeTimeout:   10 * time.Second,
-					ResponseHeaderTimeout: 10 * time.Second,
+					Proxy:               http.ProxyFromEnvironment,
+					TLSHandshakeTimeout: 10 * time.Second,
 					DialContext: (&net.Dialer{
 						Timeout: 10 * time.Second,
 					}).DialContext,
@@ -185,24 +183,24 @@ func (c *Client) CheckForUpdate(ctx context.Context, creds Credentials) (bool, e
 
 // LongPollWait sends a signed message to a DNClient API endpoint that will block, returning only
 // if there is an action the client should take before the timeout (config updates, debug commands)
-func (c *Client) LongPollWait(ctx context.Context, creds Credentials, supportedActions []string) (string, error) {
+func (c *Client) LongPollWait(ctx context.Context, creds Credentials, supportedActions []string) (*message.LongPollWaitResponse, error) {
 	value, err := json.Marshal(message.LongPollWaitRequest{
 		SupportedActions: supportedActions,
 	})
 	if err != nil {
-		return "", fmt.Errorf("failed to marshal DNClient message: %s", err)
+		return nil, fmt.Errorf("failed to marshal DNClient message: %s", err)
 	}
 
 	respBody, err := c.postDNClient(ctx, message.LongPollWait, value, creds.HostID, creds.Counter, creds.PrivateKey)
 	if err != nil {
-		return "", fmt.Errorf("failed to post message to dnclient api: %w", err)
+		return nil, fmt.Errorf("failed to post message to dnclient api: %w", err)
 	}
 	result := message.LongPollWaitResponseWrapper{}
 	err = json.Unmarshal(respBody, &result)
 	if err != nil {
-		return "", fmt.Errorf("failed to interpret API response: %s", err)
+		return nil, fmt.Errorf("failed to interpret API response: %s", err)
 	}
-	return result.Data.Action, nil
+	return &result.Data, nil
 }
 
 // DoUpdate sends a signed message to the DNClient API to fetch the new configuration update. During this call a new
@@ -282,6 +280,19 @@ func (c *Client) DoUpdate(ctx context.Context, creds Credentials) ([]byte, []byt
 	return result.Config, dhPrivkeyPEM, newCreds, nil
 }
 
+func (c *Client) CommandResponse(ctx context.Context, creds Credentials, responseToken string, response any) error {
+	value, err := json.Marshal(message.CommandResponseRequest{
+		ResponseToken: responseToken,
+		Response:      response,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to marshal DNClient message: %s", err)
+	}
+
+	_, err = c.postDNClient(ctx, message.CommandResponse, value, creds.HostID, creds.Counter, creds.PrivateKey)
+	return err
+}
+
 func (c *Client) StreamCommandResponse(ctx context.Context, creds Credentials, responseToken string) (*StreamController, error) {
 	value, err := json.Marshal(message.CommandResponseRequest{
 		ResponseToken: responseToken,
@@ -313,13 +324,11 @@ func (c *Client) streamingPostDNClient(ctx context.Context, reqType string, valu
 	sc := &StreamController{w: pw, done: done}
 
 	go func() {
-		defer func() {
-			close(done)
-		}()
+		defer close(done)
 
 		resp, err := c.streamingClient.Do(req)
 		if err != nil {
-			sc.err.Store(fmt.Errorf("failed to call dnclient endpoint: %s", err))
+			sc.err.Store(fmt.Errorf("failed to call dnclient endpoint: %w", err))
 			return
 		}
 		defer resp.Body.Close()
@@ -362,7 +371,7 @@ func (c *Client) postDNClient(ctx context.Context, reqType string, value []byte,
 	}
 	resp, err := c.client.Do(req)
 	if err != nil {
-		return nil, fmt.Errorf("failed to call dnclient endpoint: %s", err)
+		return nil, fmt.Errorf("failed to call dnclient endpoint: %w", err)
 	}
 	defer resp.Body.Close()
 
@@ -409,14 +418,18 @@ func (sc *StreamController) Err() error {
 	return err.(error)
 }
 
-// Write implements the io.Writer interface for StreamController. It writes to the request body. It never returns an
-// error. If the StreamController has already encountered an error, Write will return immediately without writing.
-// To check for errors, call Err.
-func (sc *StreamController) Write(p []byte) (n int, err error) {
+// Write implements the io.Writer interface for StreamController. It writes to the request body. If the StreamController
+// has already encountered an error, it will be returned and nothing will be written.
+func (sc *StreamController) Write(p []byte) (int, error) {
 	if sc.Err() != nil {
 		return 0, sc.Err()
 	}
-	return sc.w.Write(p)
+
+	n, err := sc.w.Write(p)
+	if err != nil {
+		sc.err.Store(err)
+	}
+	return n, err
 }
 
 // Close closes the StreamController, signaling that the request body is complete and the response can be read.

client_test.go

@@ -317,6 +317,97 @@ func TestDoUpdate(t *testing.T) {
 
 }
 
+func TestCommandResponse(t *testing.T) {
+	t.Parallel()
+
+	useragent := "testClient"
+	ts := dnapitest.NewServer(useragent)
+	t.Cleanup(func() { ts.Close() })
+
+	ca, _ := dnapitest.NebulaCACert()
+	caPEM, err := ca.MarshalToPEM()
+	require.NoError(t, err)
+
+	c := NewClient(useragent, ts.URL)
+
+	code := "foobar"
+	ts.ExpectEnrollment(code, func(req message.EnrollRequest) []byte {
+		cfg, err := yaml.Marshal(m{
+			// we need to send this or we'll get an error from the api client
+			"pki": m{"ca": string(caPEM)},
+			// here we reflect values back to the client for test purposes
+			"test": m{"code": req.Code, "dhPubkey": req.DHPubkey},
+		})
+		if err != nil {
+			return jsonMarshal(message.EnrollResponse{
+				Errors: message.APIErrors{{
+					Code:    "ERR_FAILED_TO_MARSHAL_YAML",
+					Message: "failed to marshal test response config",
+				}},
+			})
+		}
+
+		return jsonMarshal(message.EnrollResponse{
+			Data: message.EnrollResponseData{
+				HostID:      "foobar",
+				Counter:     1,
+				Config:      cfg,
+				TrustedKeys: cert.MarshalEd25519PublicKey(ca.Details.PublicKey),
+				Organization: message.EnrollResponseDataOrg{
+					ID:   "foobaz",
+					Name: "foobar's foo org",
+				},
+			},
+		})
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+	config, pkey, creds, _, err := c.Enroll(ctx, testutil.NewTestLogger(), "foobar")
+	require.NoError(t, err)
+
+	// make sure all credential values were set
+	assert.NotEmpty(t, creds.HostID)
+	assert.NotEmpty(t, creds.PrivateKey)
+	assert.NotEmpty(t, creds.TrustedKeys)
+	assert.NotEmpty(t, creds.Counter)
+
+	// make sure we got a config back
+	assert.NotEmpty(t, config)
+	assert.NotEmpty(t, pkey)
+
+	// This time sign the response with the correct CA key.
+	responseToken := "abc123"
+	res := map[string]any{"msg": "Hello, world!"}
+	ts.ExpectRequest(message.CommandResponse, http.StatusOK, func(r message.RequestWrapper) []byte {
+		var val map[string]any
+		err := json.Unmarshal(r.Value, &val)
+		require.NoError(t, err)
+		require.Contains(t, val, "responseToken")
+		require.Equal(t, responseToken, val["responseToken"])
+		require.Contains(t, val, "response")
+		require.Equal(t, res, val["response"])
+		return jsonMarshal(struct{}{})
+	})
+
+	err = c.CommandResponse(context.Background(), *creds, responseToken, res)
+	require.NoError(t, err)
+
+	// Test error handling
+	errorMsg := "sample error"
+	ts.ExpectRequest(message.CommandResponse, http.StatusBadRequest, func(r message.RequestWrapper) []byte {
+		return jsonMarshal(message.EnrollResponse{
+			Errors: message.APIErrors{{
+				Code:    "ERR_INVALID_VALUE",
+				Message: errorMsg,
+			}},
+		})
+	})
+
+	err = c.CommandResponse(context.Background(), *creds, "responseToken", map[string]any{"msg": "Hello, world!"})
+	require.Error(t, err)
+}
+
 func TestStreamCommandResponse(t *testing.T) {
 	t.Parallel()
 
@@ -451,7 +542,7 @@ func TestTimeout(t *testing.T) {
 	useragent := "TestTimeout agent"
 	c := NewClient(useragent, ts.URL)
 	// The default timeout is 1 minutes. Assert the default value.
-	assert.Equal(t, c.client.Timeout, 1*time.Minute)
+	assert.Equal(t, c.client.Timeout, 2*time.Minute)
 	// The default streaming timeout is 15 minutes. Assert the default value.
 	assert.Equal(t, c.streamingClient.Timeout, 15*time.Minute)
 	// Overwrite the default value with a 10 millisecond timeout for test brevity.

dnapitest/dnapitest.go

@@ -156,7 +156,7 @@ func (s *Server) handlerDNClient(w http.ResponseWriter, r *http.Request) {
 	// Require the expected request type, otherwise we have derailed.
 	if msg.Type != res.expectedType {
 		s.errors = append(s.errors, fmt.Errorf("%s is not expected message type %s", msg.Type, res.expectedType))
-		http.Error(w, "unexpected message type", http.StatusInternalServerError)
+		http.Error(w, fmt.Sprintf("unexpected message type %s, wanted %s", msg.Type, res.expectedType), http.StatusInternalServerError)
 		return
 	}
 

message/message.go

@@ -1,6 +1,7 @@
 package message
 
 import (
+	"encoding/json"
 	"errors"
 	"strings"
 	"time"
@@ -87,7 +88,8 @@ type LongPollWaitRequest struct {
 
 // LongPollWaitResponse is the response message associated with a LongPollWait call.
 type LongPollWaitResponse struct {
-	Action string `json:"action"` // e.g. NoOp, StreamLogs, DoUpdate
+	Action        json.RawMessage `json:"action"` // e.g. NoOp, StreamLogs, DoUpdate
+	ResponseToken string          `json:"responseToken"`
 }
 
 // CommandResponseResponseWrapper contains a response to CommandResponse inside "data."
@@ -98,6 +100,7 @@ type CommandResponseResponseWrapper struct {
 // CommandResponseRequest is the request message associated with a CommandResponse call.
 type CommandResponseRequest struct {
 	ResponseToken string `json:"responseToken"`
+	Response      any    `json:"response"`
 }
 
 // DNClientCommandResponseResponse is the response message associated with a CommandResponse call.