added ip manipulations

.github/workflows/devopstest.yml

@@ -10,7 +10,7 @@ jobs:
 
   devopstest:
     runs-on: ubuntu-latest
-    container: golang:1.17
+    container: golang:1.18
 
     services:
       postgres:

.github/workflows/statictest.yml

@@ -10,7 +10,7 @@ jobs:
 
   statictest:
     runs-on: ubuntu-latest
-    container: golang:1.17
+    container: golang:1.18
     steps:
       - name: Checkout code
         uses: actions/checkout@v2

cmd/agent/main.go

@@ -106,6 +106,7 @@ catchQuitORerror:
 	for {
 		select {
 		case <-quit:
+			cancel()
 			break catchQuitORerror
 		case err = <-er:
 			return err

cmd/server/main.go

@@ -61,7 +61,8 @@ func main() {
 	quit := make(chan os.Signal)
 	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
 
-	//See example here: https://pkg.go.dev/net/http#example-Server.Shutdown
+	// See example here: https://pkg.go.dev/net/http#example-Server.Shutdown
+	// Gracefully Shutdown
 	go func() {
 		sig := <-quit
 		logger.Info(fmt.Sprintf("caught sig: %+v", sig))

internal/agent/agent.go

@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
+	"net"
 	"net/http"
 	"os"
 	"runtime"
@@ -114,6 +115,8 @@ func sendJSONData(ctx context.Context, cfg config.Config, client *http.Client, m
 	request, err := http.NewRequest(http.MethodPost, url, &buf)
 	request.Header.Add("Content-Type", "application/json")
 	request.Header.Add("Content-Encoding", "gzip")
+	request.Header.Add("X-Real-IP", GetOutboundIP().String())
+
 	if cfg.CryptoKey.E != 0 {
 		request.Header.Add("Content-Encoding", "64base")
 	}
@@ -249,3 +252,17 @@ func StartProfiling(ctx chan int, file string, what string) {
 	}
 	fmt.Println("collecting memory really done")
 }
+
+// Get preferred outbound ip of this machine
+// https://stackoverflow.com/questions/23558425/how-do-i-get-the-local-ip-address-in-go
+func GetOutboundIP() net.IP {
+	conn, err := net.Dial("udp", "8.8.8.8:80")
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer conn.Close()
+
+	localAddr := conn.LocalAddr().(*net.UDPAddr)
+
+	return localAddr.IP
+}

internal/handlers/router.go

@@ -21,12 +21,14 @@ func MetricRouter(db storage.Repositories, cfg *config.Config, logger *zap.Logge
 	r := chi.NewRouter()
 	mh := NewMetricHandler(db, logger)
 	rsaMW := NewRsaMW(rsa.PrivateKey(cfg.CryptoKey))
+	mu := newMetricUtils(cfg)
 
 	// use inbuild middleware
 	r.Use(middleware.RequestID)
 	r.Use(middleware.RealIP)
 	r.Use(middleware.Logger)
 	r.Use(middleware.Recoverer)
+	r.Use(mu.checkTrusted)
 
 	r.Route("/update/", func(r chi.Router) {
 		r.Post("/gauge/*", Conveyor(mh.PostHandlerGouge(dbUpdated), checkForPost, checkForLength, unpackGZIP))
@@ -39,7 +41,7 @@ func MetricRouter(db storage.Repositories, cfg *config.Config, logger *zap.Logge
 	})
 
 	r.Route("/value/", func(r chi.Router) {
-		r.Get("/{type}/{name}", mh.GetHandlerValue())
+		r.Get("/{type}/{name}", Conveyor(mh.GetHandlerValue(), checkForGet))
 		r.Post("/", Conveyor(mh.PostHandlerReturn(&cfg.Key), checkForJSON, checkForPost, packGZIP, unpackGZIP))
 	})
 

internal/handlers/utils.go

@@ -6,10 +6,14 @@ import (
 	"crypto/rand"
 	"crypto/rsa"
 	"crypto/sha256"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"net/http"
+	"net/netip"
 	"strings"
+
+	"github.com/maffka123/metricCollector/internal/server/config"
 )
 
 type Middleware func(http.Handler) http.HandlerFunc
@@ -18,6 +22,14 @@ type gzipWriter struct {
 	Writer io.Writer
 }
 
+type metricUtils struct {
+	cfg *config.Config
+}
+
+func newMetricUtils(cfg *config.Config) metricUtils {
+	return metricUtils{cfg: cfg}
+}
+
 func (w gzipWriter) Write(b []byte) (int, error) {
 	return w.Writer.Write(b)
 }
@@ -43,6 +55,16 @@ func checkForPost(next http.Handler) http.HandlerFunc {
 	})
 }
 
+func checkForGet(next http.Handler) http.HandlerFunc {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodGet {
+			http.Error(w, "Only GET requests are allowed!", http.StatusMethodNotAllowed)
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}
+
 func checkForJSON(next http.Handler) http.HandlerFunc {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.Header.Get("Content-Type") != "application/json" {
@@ -142,3 +164,38 @@ func Conveyor(h http.HandlerFunc, middlewares ...Middleware) http.HandlerFunc {
 	}
 	return h
 }
+
+func (mu *metricUtils) checkTrusted(next http.Handler) http.Handler {
+	fn := func(w http.ResponseWriter, r *http.Request) {
+		if mu.cfg.TrustedSubnet != "" {
+			b, err := ifIPinCIDR(mu.cfg.TrustedSubnet, r.Header.Get("X-Real-IP"))
+			if err != nil {
+				http.Error(w, fmt.Sprintf("IP address could not be parsed: %s", err), http.StatusForbidden)
+				return
+			}
+
+			if !*b {
+				http.Error(w, "IP address is not inside relible network", http.StatusForbidden)
+				return
+			}
+		}
+		next.ServeHTTP(w, r)
+	}
+
+	return http.HandlerFunc(fn)
+}
+
+func ifIPinCIDR(cidr string, ipStr string) (*bool, error) {
+	network, err := netip.ParsePrefix(cidr)
+	if err != nil {
+		return nil, err
+	}
+
+	ip, err := netip.ParseAddr(ipStr)
+	if err != nil {
+		return nil, err
+	}
+
+	b := network.Contains(ip)
+	return &b, nil
+}

internal/handlers/utils_test.go

@@ -17,6 +17,7 @@ import (
 	"os"
 	"testing"
 
+	"github.com/maffka123/metricCollector/internal/server/config"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -117,6 +118,113 @@ func Test_checkForPost(t *testing.T) {
 	}
 }
 
+func Test_checkForGet(t *testing.T) {
+	type args struct {
+		next http.Handler
+	}
+	type want struct {
+		status int
+		body   string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    want
+		request *http.Request
+	}{
+		{name: "no post", args: args{
+			next: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }),
+		},
+			want:    want{status: 405, body: "Only GET requests are allowed!\n"},
+			request: httptest.NewRequest(http.MethodPost, "/update/", nil)},
+		{name: "post", args: args{
+			next: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }),
+		},
+			want:    want{status: 200, body: ""},
+			request: httptest.NewRequest(http.MethodGet, "/update/", nil)},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			request := tt.request
+			w := httptest.NewRecorder()
+			h := http.HandlerFunc(checkForGet(tt.args.next))
+			h.ServeHTTP(w, request)
+			result := w.Result()
+			defer result.Body.Close()
+
+			bodyBytes, err := io.ReadAll(result.Body)
+			if err != nil {
+				fmt.Println(err)
+			}
+			bodyString := string(bodyBytes)
+
+			assert.Equal(t, tt.want.status, result.StatusCode)
+			assert.Equal(t, tt.want.body, bodyString)
+		})
+	}
+}
+
+func Test_checkTrusted(t *testing.T) {
+	type args struct {
+		next   http.Handler
+		mu     metricUtils
+		header string
+	}
+	type want struct {
+		status int
+		body   string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    want
+		request *http.Request
+	}{
+		{name: "success", args: args{
+			next:   http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }),
+			mu:     metricUtils{cfg: &config.Config{TrustedSubnet: "172.17.0.0/16"}},
+			header: "172.17.0.2",
+		},
+			want:    want{status: 200, body: ""},
+			request: httptest.NewRequest(http.MethodPost, "/update/", nil)},
+
+		{name: "forbidden", args: args{
+			next:   http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }),
+			mu:     metricUtils{cfg: &config.Config{TrustedSubnet: "172.17.0.0/16"}},
+			header: "172.16.0.2",
+		},
+			want:    want{status: 403, body: "IP address is not inside relible network\n"},
+			request: httptest.NewRequest(http.MethodPost, "/update/", nil)},
+		{name: "ignore", args: args{
+			next:   http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) }),
+			mu:     metricUtils{cfg: &config.Config{}},
+			header: "172.16.0.2",
+		},
+			want:    want{status: 200, body: ""},
+			request: httptest.NewRequest(http.MethodPost, "/update/", nil)},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			request := tt.request
+			request.Header.Add("X-Real-IP", tt.args.header)
+			w := httptest.NewRecorder()
+			h := tt.args.mu.checkTrusted(tt.args.next)
+			h.ServeHTTP(w, request)
+			result := w.Result()
+			defer result.Body.Close()
+
+			bodyBytes, err := io.ReadAll(result.Body)
+			if err != nil {
+				fmt.Println(err)
+			}
+			bodyString := string(bodyBytes)
+
+			assert.Equal(t, tt.want.status, result.StatusCode)
+			assert.Equal(t, tt.want.body, bodyString)
+		})
+	}
+}
+
 func Test_unpackGZIP(t *testing.T) {
 	type args struct {
 		next http.Handler

internal/server/config/config.go

@@ -30,6 +30,7 @@ type Config struct {
 	Debug         bool          `env:"METRIC_SERVER_DEBUG"`
 	CryptoKey     rsaPrivKey    `env:"CRYPTO_KEY" json:"crypto_key"`
 	configFile    string        `env:"CONFIG"`
+	TrustedSubnet string        `env:"TRUSTED_SUBNET" json:"trusted_subnet"`
 }
 
 func (v rsaPrivKey) String() string {
@@ -67,6 +68,7 @@ func InitConfig() (Config, error) {
 	flag.Var(&cfg.CryptoKey, "ck", "crypto key for asymmetric encoding")
 	flag.BoolVar(&cfg.Debug, "debug", true, "key for hash function")
 	flag.StringVar(&cfg.configFile, "c", "", "location of config.json file")
+	flag.StringVar(&cfg.configFile, "t", "", "ip of a trusted agent")
 
 	// find full options link here: https://github.com/jackc/pgx/blob/master/pgxpool/pool.go
 	flag.StringVar(&cfg.DBpath, "d", "", "path for connection with pg: postgres://postgres:pass@localhost:5432/test?pool_max_conns=10")