feat(x): Support headers in the fetch tool (#237)

x/examples/fetch/main.go

@@ -15,13 +15,15 @@
 package main
 
 import (
+	"bufio"
 	"context"
 	"flag"
 	"fmt"
 	"io"
 	"log"
 	"net"
 	"net/http"
+	"net/textproto"
 	"os"
 	"path"
 	"strings"
@@ -32,6 +34,17 @@ import (
 
 var debugLog log.Logger = *log.New(io.Discard, "", 0)
 
+type stringArrayFlagValue []string
+
+func (v *stringArrayFlagValue) String() string {
+	return fmt.Sprint(*v)
+}
+
+func (v *stringArrayFlagValue) Set(value string) error {
+	*v = append(*v, value)
+	return nil
+}
+
 func init() {
 	flag.Usage = func() {
 		fmt.Fprintf(flag.CommandLine.Output(), "Usage: %s [flags...] <url>\n", path.Base(os.Args[0]))
@@ -44,6 +57,8 @@ func main() {
 	transportFlag := flag.String("transport", "", "Transport config")
 	addressFlag := flag.String("address", "", "Address to connect to. If empty, use the URL authority")
 	methodFlag := flag.String("method", "GET", "The HTTP method to use")
+	var headersFlag stringArrayFlagValue
+	flag.Var(&headersFlag, "H", "Raw HTTP Header line to add. It must not end in \\r\\n")
 	timeoutSecFlag := flag.Int("timeout", 5, "Timeout in seconds")
 
 	flag.Parse()
@@ -89,12 +104,28 @@ func main() {
 		}
 		return dialer.DialStream(ctx, net.JoinHostPort(host, port))
 	}
-	httpClient := &http.Client{Transport: &http.Transport{DialContext: dialContext}, Timeout: time.Duration(*timeoutSecFlag) * time.Second}
+	httpClient := &http.Client{
+		Transport: &http.Transport{DialContext: dialContext},
+		Timeout:   time.Duration(*timeoutSecFlag) * time.Second,
+		CheckRedirect: func(req *http.Request, via []*http.Request) error {
+			return http.ErrUseLastResponse
+		},
+	}
 
 	req, err := http.NewRequest(*methodFlag, url, nil)
 	if err != nil {
 		log.Fatalln("Failed to create request:", err)
 	}
+	headerText := strings.Join(headersFlag, "\r\n") + "\r\n\r\n"
+	h, err := textproto.NewReader(bufio.NewReader(strings.NewReader(headerText))).ReadMIMEHeader()
+	if err != nil {
+		log.Fatalf("invalid header line: %v", err)
+	}
+	for name, values := range h {
+		for _, value := range values {
+			req.Header.Add(name, value)
+		}
+	}
 	resp, err := httpClient.Do(req)
 	if err != nil {
 		log.Fatalf("HTTP request failed: %v\n", err)