Implement POST for gitlabnet

7 files changed, 127 insertions, 26 deletions

diff --git a/go/internal/command/twofactorrecover/twofactorrecover_test.go b/go/internal/command/twofactorrecover/twofactorrecover_test.go
index 1ab65cf..4c6f566 100644
--- a/go/internal/command/twofactorrecover/twofactorrecover_test.go
+++ b/go/internal/command/twofactorrecover/twofactorrecover_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"net/http"
 	"testing"
 
@@ -22,22 +23,31 @@ var (
 		{
 			Path: "/api/v4/internal/two_factor_recovery_codes",
 			Handler: func(w http.ResponseWriter, r *http.Request) {
-				if r.URL.Query().Get("key_id") == "1" {
+				b, _ := ioutil.ReadAll(r.Body)
+				defer r.Body.Close()
+
+				var bodyRaw map[string]*json.RawMessage
+				json.Unmarshal(b, &bodyRaw)
+
+				var keyId string
+				json.Unmarshal(*bodyRaw["key_id"], &keyId)
+
+				switch keyId {
+				case "1":
 					body := map[string]interface{}{
 						"success":        true,
 						"recovery_codes": [2]string{"recovery", "codes"},
 					}
 					json.NewEncoder(w).Encode(body)
-				} else if r.URL.Query().Get("key_id") == "broken_message" {
+				case "broken_message":
 					body := map[string]interface{}{
 						"success": false,
 						"message": "Forbidden!",
 					}
-					w.WriteHeader(http.StatusForbidden)
 					json.NewEncoder(w).Encode(body)
-				} else if r.URL.Query().Get("key_id") == "broken" {
+				case "broken":
 					w.WriteHeader(http.StatusInternalServerError)
-				} else {
+				default:
 					fmt.Fprint(w, "null")
 				}
 			},
diff --git a/go/internal/gitlabnet/client.go b/go/internal/gitlabnet/client.go
index abc218f..c2453e5 100644
--- a/go/internal/gitlabnet/client.go
+++ b/go/internal/gitlabnet/client.go
@@ -17,8 +17,7 @@ const (
 
 type GitlabClient interface {
 	Get(path string) (*http.Response, error)
-	// TODO: implement posts
-	// Post(path string) (http.Response, error)
+	Post(path string, data interface{}) (*http.Response, error)
 }
 
 type ErrorResponse struct {
diff --git a/go/internal/gitlabnet/client_test.go b/go/internal/gitlabnet/client_test.go
index f69f284..064dd5d 100644
--- a/go/internal/gitlabnet/client_test.go
+++ b/go/internal/gitlabnet/client_test.go
@@ -23,6 +23,15 @@ func TestClients(t *testing.T) {
 			},
 		},
 		{
+			Path: "/api/v4/internal/post_endpoint",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				b, _ := ioutil.ReadAll(r.Body)
+				defer r.Body.Close()
+
+				fmt.Fprint(w, "Echo: "+string(b))
+			},
+		},
+		{
 			Path: "/api/v4/internal/auth",
 			Handler: func(w http.ResponseWriter, r *http.Request) {
 				fmt.Fprint(w, r.Header.Get(secretHeaderName))
@@ -68,6 +77,7 @@ func TestClients(t *testing.T) {
 
 			testBrokenRequest(t, tc.client)
 			testSuccessfulGet(t, tc.client)
+			testSuccessfulPost(t, tc.client)
 			testMissing(t, tc.client)
 			testErrorMessage(t, tc.client)
 			testAuthenticationHeader(t, tc.client)
@@ -89,32 +99,66 @@ func testSuccessfulGet(t *testing.T, client GitlabClient) {
 	})
 }
 
+func testSuccessfulPost(t *testing.T, client GitlabClient) {
+	t.Run("Successful Post", func(t *testing.T) {
+		data := map[string]string{"key": "value"}
+
+		response, err := client.Post("/post_endpoint", data)
+		defer response.Body.Close()
+
+		require.NoError(t, err)
+		require.NotNil(t, response)
+
+		responseBody, err := ioutil.ReadAll(response.Body)
+		assert.NoError(t, err)
+		assert.Equal(t, string(responseBody), "Echo: {\"key\":\"value\"}\n")
+	})
+}
+
 func testMissing(t *testing.T, client GitlabClient) {
-	t.Run("Missing error", func(t *testing.T) {
+	t.Run("Missing error for GET", func(t *testing.T) {
 		response, err := client.Get("/missing")
 		assert.EqualError(t, err, "Internal API error (404)")
 		assert.Nil(t, response)
 	})
+
+	t.Run("Missing error for POST", func(t *testing.T) {
+		response, err := client.Post("/missing", map[string]string{})
+		assert.EqualError(t, err, "Internal API error (404)")
+		assert.Nil(t, response)
+	})
 }
 
 func testErrorMessage(t *testing.T, client GitlabClient) {
-	t.Run("Error with message", func(t *testing.T) {
+	t.Run("Error with message for GET", func(t *testing.T) {
 		response, err := client.Get("/error")
 		assert.EqualError(t, err, "Don't do that")
 		assert.Nil(t, response)
 	})
+
+	t.Run("Error with message for POST", func(t *testing.T) {
+		response, err := client.Post("/error", map[string]string{})
+		assert.EqualError(t, err, "Don't do that")
+		assert.Nil(t, response)
+	})
 }
 
 func testBrokenRequest(t *testing.T, client GitlabClient) {
-	t.Run("Broken request", func(t *testing.T) {
+	t.Run("Broken request for GET", func(t *testing.T) {
 		response, err := client.Get("/broken")
 		assert.EqualError(t, err, "Internal API unreachable")
 		assert.Nil(t, response)
 	})
+
+	t.Run("Broken request for POST", func(t *testing.T) {
+		response, err := client.Post("/broken", map[string]string{})
+		assert.EqualError(t, err, "Internal API unreachable")
+		assert.Nil(t, response)
+	})
 }
 
 func testAuthenticationHeader(t *testing.T, client GitlabClient) {
-	t.Run("Authentication headers", func(t *testing.T) {
+	t.Run("Authentication headers for GET", func(t *testing.T) {
 		response, err := client.Get("/auth")
 		defer response.Body.Close()
 
@@ -128,4 +172,19 @@ func testAuthenticationHeader(t *testing.T, client GitlabClient) {
 		require.NoError(t, err)
 		assert.Equal(t, "sssh, it's a secret", string(header))
 	})
+
+	t.Run("Authentication headers for POST", func(t *testing.T) {
+		response, err := client.Post("/auth", map[string]string{})
+		defer response.Body.Close()
+
+		require.NoError(t, err)
+		require.NotNil(t, response)
+
+		responseBody, err := ioutil.ReadAll(response.Body)
+		require.NoError(t, err)
+
+		header, err := base64.StdEncoding.DecodeString(string(responseBody))
+		require.NoError(t, err)
+		assert.Equal(t, "sssh, it's a secret", string(header))
+	})
 }
diff --git a/go/internal/gitlabnet/socketclient.go b/go/internal/gitlabnet/socketclient.go
index 3bd7c70..64c3005 100644
--- a/go/internal/gitlabnet/socketclient.go
+++ b/go/internal/gitlabnet/socketclient.go
@@ -1,7 +1,9 @@
 package gitlabnet
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"net"
 	"net/http"
 	"strings"
@@ -44,3 +46,21 @@ func (c *GitlabSocketClient) Get(path string) (*http.Response, error) {
 
 	return doRequest(c.httpClient, c.config, request)
 }
+
+func (c *GitlabSocketClient) Post(path string, data interface{}) (*http.Response, error) {
+	path = normalizePath(path)
+
+	buffer := new(bytes.Buffer)
+	if err := json.NewEncoder(buffer).Encode(data); err != nil {
+		return nil, err
+	}
+
+	request, err := http.NewRequest("POST", socketBaseUrl+path, buffer)
+	request.Header.Add("Content-Type", "application/json")
+
+	if err != nil {
+		return nil, err
+	}
+
+	return doRequest(c.httpClient, c.config, request)
+}
diff --git a/go/internal/gitlabnet/twofactorrecover/client.go b/go/internal/gitlabnet/twofactorrecover/client.go
index bce0fcb..cbec6c1 100644
--- a/go/internal/gitlabnet/twofactorrecover/client.go
+++ b/go/internal/gitlabnet/twofactorrecover/client.go
@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
-	"net/url"
 
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet"
@@ -22,6 +21,10 @@ type Response struct {
 	Message       string   `json:"message"`
 }
 
+type RequestBody struct {
+	KeyId string `json:"key_id"`
+}
+
 func NewClient(config *config.Config) (*Client, error) {
 	client, err := gitlabnet.GetClient(config)
 	if err != nil {
@@ -32,11 +35,8 @@ func NewClient(config *config.Config) (*Client, error) {
 }
 
 func (c *Client) GetRecoveryCodes(gitlabKeyId string) ([]string, error) {
-	values := url.Values{}
-	values.Add("key_id", gitlabKeyId)
-
-	path := "/two_factor_recovery_codes?" + values.Encode()
-	response, err := c.client.Get(path)
+	values := RequestBody{KeyId: gitlabKeyId}
+	response, err := c.client.Post("/two_factor_recovery_codes", values)
 
 	if err != nil {
 		return nil, err
diff --git a/go/internal/gitlabnet/twofactorrecover/client_test.go b/go/internal/gitlabnet/twofactorrecover/client_test.go
index 094b029..79980c4 100644
--- a/go/internal/gitlabnet/twofactorrecover/client_test.go
+++ b/go/internal/gitlabnet/twofactorrecover/client_test.go
@@ -3,6 +3,7 @@ package twofactorrecover
 import (
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"net/http"
 	"testing"
 
@@ -25,29 +26,39 @@ func init() {
 		{
 			Path: "/api/v4/internal/two_factor_recovery_codes",
 			Handler: func(w http.ResponseWriter, r *http.Request) {
-				if r.URL.Query().Get("key_id") == "1" {
+				b, _ := ioutil.ReadAll(r.Body)
+				defer r.Body.Close()
+
+				var bodyRaw map[string]*json.RawMessage
+				json.Unmarshal(b, &bodyRaw)
+
+				var keyId string
+				json.Unmarshal(*bodyRaw["key_id"], &keyId)
+
+				switch keyId {
+				case "0":
 					body := map[string]interface{}{
 						"success":        true,
 						"recovery_codes": [2]string{"recovery 1", "codes 1"},
 					}
 					json.NewEncoder(w).Encode(body)
-				} else if r.URL.Query().Get("key_id") == "0" {
+				case "1":
 					body := map[string]interface{}{
 						"success": false,
 						"message": "missing user",
 					}
 					json.NewEncoder(w).Encode(body)
-				} else if r.URL.Query().Get("key_id") == "2" {
+				case "2":
 					w.WriteHeader(http.StatusForbidden)
 					body := &gitlabnet.ErrorResponse{
 						Message: "Not allowed!",
 					}
 					json.NewEncoder(w).Encode(body)
-				} else if r.URL.Query().Get("key_id") == "3" {
+				case "3":
 					w.Write([]byte("{ \"message\": \"broken json!\""))
-				} else if r.URL.Query().Get("key_id") == "4" {
+				case "4":
 					w.WriteHeader(http.StatusForbidden)
-				} else {
+				default:
 					fmt.Fprint(w, "null")
 				}
 			},
@@ -59,7 +70,7 @@ func TestGetRecoveryCodes(t *testing.T) {
 	client, cleanup := setup(t)
 	defer cleanup()
 
-	result, err := client.GetRecoveryCodes("1")
+	result, err := client.GetRecoveryCodes("0")
 	assert.NoError(t, err)
 	assert.Equal(t, []string{"recovery 1", "codes 1"}, result)
 }
@@ -68,7 +79,7 @@ func TestMissingUser(t *testing.T) {
 	client, cleanup := setup(t)
 	defer cleanup()
 
-	_, err := client.GetRecoveryCodes("0")
+	_, err := client.GetRecoveryCodes("1")
 	assert.Equal(t, "missing user", err.Error())
 }
 
diff --git a/spec/gitlab_shell_two_factor_recovery_spec.rb b/spec/gitlab_shell_two_factor_recovery_spec.rb
index 719d57e..872fa85 100644
--- a/spec/gitlab_shell_two_factor_recovery_spec.rb
+++ b/spec/gitlab_shell_two_factor_recovery_spec.rb
@@ -10,7 +10,9 @@ describe 'bin/gitlab-shell 2fa_recovery_codes' do
       res.content_type = 'application/json'
       res.status = 200
 
-      if req.query['key_id'] == '100'
+      key_id = req.query['key_id'] || JSON.parse(req.body)['key_id']
+
+      if key_id == '100'
         res.body = '{"success":true, "recovery_codes": ["1", "2"]}'
       else
         res.body = '{"success":false, "message": "Forbidden!"}'