From 2b565e8f12888bb262417db4f295a5f86bcaf6a6 Mon Sep 17 00:00:00 2001
From: zhangnew <zhangneww@gmail.com>
Date: Mon, 16 Oct 2023 19:42:58 +0800
Subject: [PATCH] HttpGet support Alt-Svc (#3)

* HttpGet support Alt-Svc

* add version params

* fix code review
---
 README.md |  3 +++
 go.mod    |  1 +
 go.sum    |  2 ++
 main.go   | 70 +++++++++++++++++++++++++++++++++++++++++++++++++++++--
 4 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index f0b8ce9..dff8a05 100644
--- a/README.md
+++ b/README.md
@@ -6,4 +6,7 @@ Agent of Nezha Monitoring
 
 <!--GAMFC_DELIMITER--><a href="https://github.com/naiba" title="naiba">
   <img src="https://avatars.githubusercontent.com/u/29243953?v=4" width="50;" alt="naiba"/>
+</a>
+<a href="https://github.com/zhangnew" title="zhangnew">
+  <img src="https://avatars.githubusercontent.com/u/9146834?v=4" width="50;" alt="zhangnew"/>
 </a><!--GAMFC_DELIMITER_END-->
diff --git a/go.mod b/go.mod
index 28ce7bd..2bb8c3f 100644
--- a/go.mod
+++ b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/blang/semver v3.5.1+incompatible
 	github.com/creack/pty v1.1.18
 	github.com/dean2021/goss v0.0.0-20230129073947-df90431348f1
+	github.com/ebi-yade/altsvc-go v0.1.1
 	github.com/go-ping/ping v1.1.0
 	github.com/gorilla/websocket v1.5.0
 	github.com/iamacarpet/go-winpty v1.0.4
diff --git a/go.sum b/go.sum
index 45ac133..39f97f8 100644
--- a/go.sum
+++ b/go.sum
@@ -70,6 +70,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dean2021/goss v0.0.0-20230129073947-df90431348f1 h1:5UiJ324LiCdOF/3w/5IeXrKVjdnwHoalvLG2smb3wi4=
 github.com/dean2021/goss v0.0.0-20230129073947-df90431348f1/go.mod h1:NiLueuVb3hYcdF4ta+2ezcKJh6BEjhrBz9Hts6XJ5Sc=
+github.com/ebi-yade/altsvc-go v0.1.1 h1:HmZDNb5ZOPlkyXhi34LnRckawFCux7yPYw+dtInIixo=
+github.com/ebi-yade/altsvc-go v0.1.1/go.mod h1:K/U20bLcsOVrbTeDhqRjp+e3tgNT5iAqSiQzPoU0/Q0=
 github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98=
diff --git a/main.go b/main.go
index a83e7af..4c946ad 100644
--- a/main.go
+++ b/main.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"net/url"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -18,6 +19,7 @@ import (
 	"github.com/AlecAivazis/survey/v2"
 	bpc "github.com/DaRealFreak/cloudflare-bp-go"
 	"github.com/blang/semver"
+	"github.com/ebi-yade/altsvc-go"
 	"github.com/go-ping/ping"
 	"github.com/gorilla/websocket"
 	"github.com/nezhahq/go-github-selfupdate/selfupdate"
@@ -48,6 +50,7 @@ type AgentCliParam struct {
 	ClientSecret          string // 客户端密钥
 	ReportDelay           int    // 报告间隔
 	TLS                   bool   // 是否使用TLS加密传输至服务端
+	Version               bool   // 当前版本号
 }
 
 var (
@@ -155,8 +158,14 @@ func main() {
 	flag.BoolVar(&agentCliParam.DisableAutoUpdate, "disable-auto-update", false, "禁用自动升级")
 	flag.BoolVar(&agentCliParam.DisableForceUpdate, "disable-force-update", false, "禁用强制升级")
 	flag.BoolVar(&agentCliParam.TLS, "tls", false, "启用SSL/TLS加密")
+	flag.BoolVarP(&agentCliParam.Version, "version", "v", false, "查看当前版本号")
 	flag.Parse()
 
+	if agentCliParam.Version {
+		fmt.Println(version)
+		return
+	}
+
 	if isEditAgentConfig {
 		editAgentConfig()
 		return
@@ -385,7 +394,12 @@ func handleIcmpPingTask(task *pb.Task, result *pb.TaskResult) {
 
 func handleHttpGetTask(task *pb.Task, result *pb.TaskResult) {
 	start := time.Now()
-	resp, err := httpClient.Get(task.GetData())
+	taskUrl := task.GetData()
+	resp, err := httpClient.Get(taskUrl)
+	checkHttpResp(taskUrl, start, resp, err, result)
+}
+
+func checkHttpResp(taskUrl string, start time.Time, resp *http.Response, err error, result *pb.TaskResult) {
 	if err == nil {
 		// 检查 HTTP Response 状态
 		result.Delay = float32(time.Since(start).Microseconds()) / 1000.0
@@ -399,13 +413,65 @@ func handleHttpGetTask(task *pb.Task, result *pb.TaskResult) {
 			c := resp.TLS.PeerCertificates[0]
 			result.Data = c.Issuer.CommonName + "|" + c.NotAfter.String()
 		}
-		result.Successful = true
+		altSvc := resp.Header.Get("Alt-Svc")
+		if altSvc != "" {
+			checkAltSvc(altSvc, taskUrl, result)
+		} else {
+			result.Successful = true
+		}
 	} else {
 		// HTTP 请求失败
 		result.Data = err.Error()
 	}
 }
 
+func checkAltSvc(altSvcStr string, taskUrl string, result *pb.TaskResult) {
+	altSvcList, err := altsvc.Parse(altSvcStr)
+	if err != nil {
+		result.Data = err.Error()
+		result.Successful = false
+		return
+	}
+
+	parsedUrl, _ := url.Parse(taskUrl)
+	originalHost := parsedUrl.Hostname()
+	originalPort := parsedUrl.Port()
+	if originalPort == "" {
+		switch parsedUrl.Scheme {
+		case "http":
+			originalPort = "80"
+		case "https":
+			originalPort = "443"
+		}
+	}
+
+	altAuthorityHost := ""
+	altAuthorityPort := ""
+	for _, altSvc := range altSvcList {
+		if altSvc.AltAuthority.Host != "" {
+			altAuthorityHost = altSvc.AltAuthority.Host
+		}
+		altAuthorityPort = altSvc.AltAuthority.Port
+	}
+	if altAuthorityHost == "" {
+		altAuthorityHost = originalHost
+	}
+	if altAuthorityHost == originalHost && altAuthorityPort == originalPort {
+		result.Successful = true
+		return
+	}
+
+	altAuthorityUrl := "https://" + altAuthorityHost + ":" + altAuthorityPort + "/"
+	req, _ := http.NewRequest("GET", altAuthorityUrl, nil)
+	req.Host = originalHost
+	req.Header.Add("Upgrade", originalHost)
+	req.Header.Add("Connection", "Upgrade")
+
+	start := time.Now()
+	resp, err := httpClient.Do(req)
+	checkHttpResp(taskUrl, start, resp, err, result)
+}
+
 func handleCommandTask(task *pb.Task, result *pb.TaskResult) {
 	if agentCliParam.DisableCommandExecute {
 		result.Data = "此 Agent 已禁止命令执行"