Rewrite bin/check in Go

.gitignore

@@ -18,3 +18,4 @@ hooks/*.d
 /bin/gitaly-upload-pack
 /bin/gitaly-receive-pack
 /bin/gitaly-upload-archive
+/bin/check

bin/check

-#!/usr/bin/env ruby
-
-require_relative '../lib/gitlab_init'
-require_relative '../lib/gitlab_net'
-
-#
-# GitLab shell check task
-#
-
-print "Check GitLab API access: "
-begin
-  resp = GitlabNet.new.check
-
-  if resp.code != "200"
-    abort "FAILED. code: #{resp.code}"
-  end
-
-  puts 'OK'
-
-  check_values = JSON.parse(resp.body)
-
-  print 'Redis available via internal API: '
-  if check_values['redis']
-    puts 'OK'
-  else
-    abort 'FAILED'
-  end
-rescue GitlabNet::ApiUnreachableError
-  abort "FAILED: Failed to connect to internal API"
-end
-
-puts "\n"

go/cmd/check/main.go

+package main
+
+import (
+	"fmt"
+	"os"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/readwriter"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
+)
+
+func main() {
+	readWriter := &readwriter.ReadWriter{
+		Out:    os.Stdout,
+		In:     os.Stdin,
+		ErrOut: os.Stderr,
+	}
+
+	executable, err := executable.New(executable.Healthcheck)
+	if err != nil {
+		fmt.Fprintln(readWriter.ErrOut, "Failed to determine executable, exiting")
+		os.Exit(1)
+	}
+
+	config, err := config.NewFromDir(executable.RootDir)
+	if err != nil {
+		fmt.Fprintln(readWriter.ErrOut, "Failed to read config, exiting")
+		os.Exit(1)
+	}
+
+	cmd, err := command.New(executable, os.Args[1:], config, readWriter)
+	if err != nil {
+		fmt.Fprintf(readWriter.ErrOut, "%v\n", err)
+		os.Exit(1)
+	}
+
+	if err = cmd.Execute(); err != nil {
+		fmt.Fprintf(readWriter.ErrOut, "%v\n", err)
+		os.Exit(1)
+	}
+}

go/internal/command/command.go

@@ -5,6 +5,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/authorizedprincipals"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/discover"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/healthcheck"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/lfsauthenticate"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/readwriter"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/receivepack"
@@ -41,6 +42,8 @@ func buildCommand(e *executable.Executable, args commandargs.CommandArgs, config
 		return buildAuthorizedKeysCommand(args.(*commandargs.AuthorizedKeys), config, readWriter)
 	case executable.AuthorizedPrincipalsCheck:
 		return buildAuthorizedPrincipalsCommand(args.(*commandargs.AuthorizedPrincipals), config, readWriter)
+	case executable.Healthcheck:
+		return buildHealthcheckCommand(config, readWriter)
 	}
 
 	return nil
@@ -72,3 +75,7 @@ func buildAuthorizedKeysCommand(args *commandargs.AuthorizedKeys, config *config
 func buildAuthorizedPrincipalsCommand(args *commandargs.AuthorizedPrincipals, config *config.Config, readWriter *readwriter.ReadWriter) Command {
 	return &authorizedprincipals.Command{Config: config, Args: args, ReadWriter: readWriter}
 }
+
+func buildHealthcheckCommand(config *config.Config, readWriter *readwriter.ReadWriter) Command {
+	return &healthcheck.Command{Config: config, ReadWriter: readWriter}
+}

go/internal/command/command_test.go

@@ -9,6 +9,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/authorizedkeys"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/authorizedprincipals"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/discover"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/healthcheck"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/lfsauthenticate"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/receivepack"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/shared/disallowedcommand"
@@ -23,6 +24,7 @@ import (
 var (
 	authorizedKeysExec       = &executable.Executable{Name: executable.AuthorizedKeysCheck}
 	authorizedPrincipalsExec = &executable.Executable{Name: executable.AuthorizedPrincipalsCheck}
+	checkExec                = &executable.Executable{Name: executable.Healthcheck}
 	gitlabShellExec          = &executable.Executable{Name: executable.GitlabShell}
 
 	basicConfig = &config.Config{GitlabUrl: "http+unix://gitlab.socket"}
@@ -35,6 +37,13 @@ func buildEnv(command string) map[string]string {
 	}
 }
 
+func buildCheckAllowedEnv(command string) map[string]string {
+	out := buildEnv(command)
+	out["GITLAB_SHELL_ALLOW_CHECK_COMMAND"] = "1"
+
+	return out
+}
+
 func TestNew(t *testing.T) {
 	testCases := []struct {
 		desc         string
@@ -79,6 +88,11 @@ func TestNew(t *testing.T) {
 			environment:  buildEnv("git-upload-archive"),
 			expectedType: &uploadarchive.Command{},
 		},
+		{
+			desc:         "it returns a Healthcheck command",
+			executable:   checkExec,
+			expectedType: &healthcheck.Command{},
+		},
 		{
 			desc:         "it returns a AuthorizedKeys command",
 			executable:   authorizedKeysExec,

go/internal/command/healthcheck/healthcheck.go

+package healthcheck
+
+import (
+	"fmt"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/readwriter"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/healthcheck"
+)
+
+var (
+	apiMessage   = "Internal API available"
+	redisMessage = "Redis available via internal API"
+)
+
+type Command struct {
+	Config     *config.Config
+	ReadWriter *readwriter.ReadWriter
+}
+
+func (c *Command) Execute() error {
+	response, err := c.runCheck()
+	if err != nil {
+		return fmt.Errorf("%v: FAILED - %v", apiMessage, err)
+	}
+
+	fmt.Fprintf(c.ReadWriter.Out, "%v: OK\n", apiMessage)
+
+	if !response.Redis {
+		return fmt.Errorf("%v: FAILED", redisMessage)
+	}
+
+	fmt.Fprintf(c.ReadWriter.Out, "%v: OK\n", redisMessage)
+	return nil
+}
+
+func (c *Command) runCheck() (*healthcheck.Response, error) {
+	client, err := healthcheck.NewClient(c.Config)
+	if err != nil {
+		return nil, err
+	}
+
+	response, err := client.Check()
+	if err != nil {
+		return nil, err
+	}
+
+	return response, nil
+}

go/internal/command/healthcheck/healthcheck_test.go

+package healthcheck
+
+import (
+	"bytes"
+	"encoding/json"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/readwriter"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/healthcheck"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/testserver"
+)
+
+var (
+	okResponse = &healthcheck.Response{
+		APIVersion:     "v4",
+		GitlabVersion:  "v12.0.0-ee",
+		GitlabRevision: "3b13818e8330f68625d80d9bf5d8049c41fbe197",
+		Redis:          true,
+	}
+
+	badRedisResponse = &healthcheck.Response{Redis: false}
+
+	okHandlers       = buildTestHandlers(200, okResponse)
+	badRedisHandlers = buildTestHandlers(200, badRedisResponse)
+	brokenHandlers   = buildTestHandlers(500, nil)
+)
+
+func buildTestHandlers(code int, rsp *healthcheck.Response) []testserver.TestRequestHandler {
+	return []testserver.TestRequestHandler{
+		{
+			Path: "/api/v4/internal/check",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(code)
+				if rsp != nil {
+					json.NewEncoder(w).Encode(rsp)
+				}
+			},
+		},
+	}
+}
+
+func TestExecute(t *testing.T) {
+	url, cleanup := testserver.StartSocketHttpServer(t, okHandlers)
+	defer cleanup()
+
+	buffer := &bytes.Buffer{}
+	cmd := &Command{
+		Config:     &config.Config{GitlabUrl: url},
+		ReadWriter: &readwriter.ReadWriter{Out: buffer},
+	}
+
+	err := cmd.Execute()
+
+	require.NoError(t, err)
+	require.Equal(t, "Internal API available: OK\nRedis available via internal API: OK\n", buffer.String())
+}
+
+func TestFailingRedisExecute(t *testing.T) {
+	url, cleanup := testserver.StartSocketHttpServer(t, badRedisHandlers)
+	defer cleanup()
+
+	buffer := &bytes.Buffer{}
+	cmd := &Command{
+		Config:     &config.Config{GitlabUrl: url},
+		ReadWriter: &readwriter.ReadWriter{Out: buffer},
+	}
+
+	err := cmd.Execute()
+	require.Error(t, err, "Redis available via internal API: FAILED")
+	require.Equal(t, "Internal API available: OK\n", buffer.String())
+}
+
+func TestFailingAPIExecute(t *testing.T) {
+	url, cleanup := testserver.StartSocketHttpServer(t, brokenHandlers)
+	defer cleanup()
+
+	buffer := &bytes.Buffer{}
+	cmd := &Command{
+		Config:     &config.Config{GitlabUrl: url},
+		ReadWriter: &readwriter.ReadWriter{Out: buffer},
+	}
+
+	err := cmd.Execute()
+	require.Empty(t, buffer.String())
+	require.EqualError(t, err, "Internal API available: FAILED - Internal API error (500)")
+}

go/internal/executable/executable.go

@@ -7,6 +7,7 @@ import (
 
 const (
 	BinDir                    = "bin"
+	Healthcheck               = "check"
 	GitlabShell               = "gitlab-shell"
 	AuthorizedKeysCheck       = "gitlab-shell-authorized-keys-check"
 	AuthorizedPrincipalsCheck = "gitlab-shell-authorized-principals-check"

go/internal/gitlabnet/healthcheck/client.go

+package healthcheck
+
+import (
+	"fmt"
+	"net/http"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet"
+)
+
+const (
+	checkPath = "/check"
+)
+
+type Client struct {
+	config *config.Config
+	client *gitlabnet.GitlabClient
+}
+
+type Response struct {
+	APIVersion     string `json:"api_version"`
+	GitlabVersion  string `json:"gitlab_version"`
+	GitlabRevision string `json:"gitlab_rev"`
+	Redis          bool   `json:"redis"`
+}
+
+func NewClient(config *config.Config) (*Client, error) {
+	client, err := gitlabnet.GetClient(config)
+	if err != nil {
+		return nil, fmt.Errorf("Error creating http client: %v", err)
+	}
+
+	return &Client{config: config, client: client}, nil
+}
+
+func (c *Client) Check() (*Response, error) {
+	resp, err := c.client.Get(checkPath)
+	if err != nil {
+		return nil, err
+	}
+
+	defer resp.Body.Close()
+
+	return parse(resp)
+}
+
+func parse(hr *http.Response) (*Response, error) {
+	response := &Response{}
+	if err := gitlabnet.ParseJSON(hr, response); err != nil {
+		return nil, err
+	}
+
+	return response, nil
+}

go/internal/gitlabnet/healthcheck/client_test.go

+package healthcheck
+
+import (
+	"encoding/json"
+	"net/http"
+	"testing"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/gitlabnet/testserver"
+
+	"github.com/stretchr/testify/require"
+)
+
+var (
+	requests = []testserver.TestRequestHandler{
+		{
+			Path: "/api/v4/internal/check",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				json.NewEncoder(w).Encode(testResponse)
+			},
+		},
+	}
+
+	testResponse = &Response{
+		APIVersion:     "v4",
+		GitlabVersion:  "v12.0.0-ee",
+		GitlabRevision: "3b13818e8330f68625d80d9bf5d8049c41fbe197",
+		Redis:          true,
+	}
+)
+
+func TestCheck(t *testing.T) {
+	client, cleanup := setup(t)
+	defer cleanup()
+
+	result, err := client.Check()
+	require.NoError(t, err)
+	require.Equal(t, testResponse, result)
+}
+
+func setup(t *testing.T) (*Client, func()) {
+	url, cleanup := testserver.StartSocketHttpServer(t, requests)
+
+	client, err := NewClient(&config.Config{GitlabUrl: url})
+	require.NoError(t, err)
+
+	return client, cleanup
+}