network: improve ctrl+c interupts

This cancels the HTTP request when Ctrl+C is made, without crashing cmk
and immediately getting the shell back to interacting users.

Fixes #126

Signed-off-by: Rohit Yadav <[email protected]>

Author: rohityadavcloud

cmd/api.go

@@ -79,7 +79,9 @@ func init() {
 
 			response, err := NewAPIRequest(r, api.Name, apiArgs, api.Async)
 			if err != nil {
-				if response != nil {
+				if strings.HasSuffix(err.Error(), "context canceled") {
+					return nil
+				} else if response != nil {
 					printResult(r.Config.Core.Output, response, nil)
 				}
 				return err

cmd/network.go

@@ -30,7 +30,6 @@ import (
 	"net/http/cookiejar"
 	"net/url"
 	"os"
-	"os/signal"
 	"sort"
 	"strings"
 	"time"
@@ -149,19 +148,11 @@ func pollAsyncJob(r *Request, jobID string) (map[string]interface{}, error) {
 	spinner := r.Config.StartSpinner("polling for async API result")
 	defer r.Config.StopSpinner(spinner)
 
-	interrupted := false
-	c := make(chan os.Signal, 1)
-	signal.Notify(c, os.Interrupt)
-	go func() {
-		<-c
-		interrupted = true
-	}()
-
 	for {
-		if interrupted {
-			return nil, errors.New("API job query interrupted")
-		}
 		select {
+		case <-r.Config.C:
+			return nil, errors.New("async API job polling interrupted")
+
 		case <-timeout.C:
 			return nil, errors.New("async API job query timed out")
 
@@ -294,10 +285,12 @@ func NewAPIRequest(r *Request, api string, args []string, isAsync bool) (map[str
 
 // we can implement further conditions to do POST or GET (or other http commands) here
 func executeRequest(r *Request, requestURL string, params url.Values) (*http.Response, error) {
+	config.SetupContext(r.Config)
 	if params.Has("password") || params.Has("userdata") {
 		requestURL = fmt.Sprintf("%s", r.Config.ActiveProfile.URL)
 		return r.Client().PostForm(requestURL, params)
 	} else {
-		return r.Client().Get(requestURL)
+		req, _ := http.NewRequestWithContext(*r.Config.Context, "GET", requestURL, nil)
+		return r.Client().Do(req)
 	}
 }

config/config.go

@@ -18,11 +18,13 @@
 package config
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"net/http"
 	"net/http/cookiejar"
 	"os"
+	"os/signal"
 	"path"
 	"path/filepath"
 	"strconv"
@@ -76,6 +78,9 @@ type Config struct {
 	HasShell      bool
 	Core          *Core
 	ActiveProfile *ServerProfile
+	Context       *context.Context
+	Cancel        context.CancelFunc
+	C             chan bool
 }
 
 func GetOutputFormats() []string {
@@ -182,7 +187,22 @@ func GetProfiles() []string {
 	return profiles
 }
 
+func SetupContext(cfg *Config) {
+	cfg.C = make(chan bool)
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, os.Interrupt)
+	ctx, cancel := context.WithCancel(context.Background())
+	cfg.Context = &ctx
+	cfg.Cancel = cancel
+	go func() {
+		<-signals
+		cfg.Cancel()
+		cfg.C <- true
+	}()
+}
+
 func newHTTPClient(cfg *Config) *http.Client {
+	SetupContext(cfg)
 	jar, _ := cookiejar.New(nil)
 	client := &http.Client{
 		Jar: jar,