From 0a9d50c054a36d5d0d9fb8e97120913cf7df9f59 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Thu, 5 May 2022 21:20:20 +0800
Subject: [PATCH] feat: print connected network info

---
 hrp/internal/httpstat/main.go | 16 +++++++++++++++-
 hrp/step_request.go           | 16 ++++++++++++++++
 hrp/step_request_test.go      | 19 ++++++++-----------
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/hrp/internal/httpstat/main.go b/hrp/internal/httpstat/main.go
index e5731898..29bf464d 100644
--- a/hrp/internal/httpstat/main.go
+++ b/hrp/internal/httpstat/main.go
@@ -97,6 +97,9 @@ type Stat struct {
 
 	// https or http
 	schema string
+
+	// connected network info
+	network, addr string
 }
 
 // Finish sets the time when reading response is done.
@@ -131,6 +134,14 @@ func (s *Stat) Durations() map[string]int64 {
 }
 
 func (s *Stat) Print() {
+	if s.network != "" && s.addr != "" {
+		printf("\n%s %s: %s\n",
+			color.CyanString("Connected to"),
+			color.YellowString(s.network),
+			color.BlueString(s.addr),
+		)
+	}
+
 	switch s.schema {
 	case "https":
 		printf(colorize(httpsTemplate),
@@ -178,7 +189,10 @@ func WithHTTPStat(req *http.Request, s *Stat) context.Context {
 			s.NameLookup = s.DNSLookup
 		},
 
-		ConnectStart: func(_, _ string) {
+		ConnectStart: func(network, addr string) {
+			s.network = network
+			s.addr = addr
+
 			s.tcpStart = time.Now()
 
 			// When connecting to IP (When no DNS lookup)
diff --git a/hrp/step_request.go b/hrp/step_request.go
index 41e5d022..30c5988e 100644
--- a/hrp/step_request.go
+++ b/hrp/step_request.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"compress/gzip"
 	"compress/zlib"
+	"crypto/tls"
 	"fmt"
 	"io"
 	"net/http"
@@ -15,6 +16,7 @@ import (
 	"time"
 
 	"github.com/andybalholm/brotli"
+	"github.com/fatih/color"
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
 
@@ -423,8 +425,22 @@ func printRequest(req *http.Request) error {
 	return nil
 }
 
+func printf(format string, a ...interface{}) (n int, err error) {
+	return fmt.Fprintf(color.Output, format, a...)
+}
+
 func printResponse(resp *http.Response) error {
 	fmt.Println("==================== response ====================")
+	connectedVia := "plaintext"
+	if resp.TLS != nil {
+		switch resp.TLS.Version {
+		case tls.VersionTLS12:
+			connectedVia = "TLSv1.2"
+		case tls.VersionTLS13:
+			connectedVia = "TLSv1.3"
+		}
+	}
+	printf("%s %s\n", color.CyanString("Connected via"), color.BlueString("%s", connectedVia))
 	respContentType := resp.Header.Get("Content-Type")
 	printBody := shouldPrintBody(respContentType)
 	respDump, err := httputil.DumpResponse(resp, printBody)
diff --git a/hrp/step_request_test.go b/hrp/step_request_test.go
index 46c62cae..1076b31d 100644
--- a/hrp/step_request_test.go
+++ b/hrp/step_request_test.go
@@ -135,34 +135,31 @@ func TestRunRequestStatOn(t *testing.T) {
 	if !assert.Greater(t, stat["Total"], int64(5)) {
 		t.Fatal()
 	}
-	if !assert.Less(t, stat["Total"]-summary.Records[0].Elapsed, int64(2)) {
+	if !assert.Less(t, stat["Total"]-summary.Records[0].Elapsed, int64(3)) {
 		t.Fatal()
 	}
 
 	// reuse connection
 	stat = summary.Records[1].HttpStat
-	if !assert.Equal(t, stat["DNSLookup"], int64(0)) {
+	if !assert.Equal(t, int64(0), stat["DNSLookup"]) {
 		t.Fatal()
 	}
-	if !assert.Equal(t, stat["TCPConnection"], int64(0)) {
+	if !assert.Equal(t, int64(0), stat["TCPConnection"]) {
 		t.Fatal()
 	}
-	if !assert.Equal(t, stat["TLSHandshake"], int64(0)) {
+	if !assert.Equal(t, int64(0), stat["TLSHandshake"]) {
 		t.Fatal()
 	}
 	if !assert.Greater(t, stat["ServerProcessing"], int64(1)) {
 		t.Fatal()
 	}
-	if !assert.Equal(t, stat["ContentTransfer"], int64(0)) {
+	if !assert.Equal(t, int64(0), stat["NameLookup"]) {
 		t.Fatal()
 	}
-	if !assert.Equal(t, stat["NameLookup"], int64(0)) {
+	if !assert.Equal(t, int64(0), stat["Connect"]) {
 		t.Fatal()
 	}
-	if !assert.Equal(t, stat["Connect"], int64(0)) {
-		t.Fatal()
-	}
-	if !assert.Equal(t, stat["Pretransfer"], int64(0)) {
+	if !assert.Equal(t, int64(0), stat["Pretransfer"]) {
 		t.Fatal()
 	}
 	if !assert.Greater(t, stat["StartTransfer"], int64(0)) {
@@ -171,7 +168,7 @@ func TestRunRequestStatOn(t *testing.T) {
 	if !assert.Greater(t, stat["Total"], int64(1)) {
 		t.Fatal()
 	}
-	if !assert.Less(t, stat["Total"]-summary.Records[0].Elapsed, int64(2)) {
+	if !assert.Less(t, stat["Total"]-summary.Records[0].Elapsed, int64(3)) {
 		t.Fatal()
 	}
 }