From b8a23e512b6b84b496dc332be66caacad9ed56b5 Mon Sep 17 00:00:00 2001
From: JustSong <quanpengsong@gmail.com>
Date: Thu, 28 Jul 2022 23:08:15 +0800
Subject: [PATCH] :bug: bug fix for client

---
 client/README.md   | 13 +++++++---
 client/main.go     | 62 ++++++++++++++++++++++++++++++----------------
 client/starter.vbs |  3 +++
 3 files changed, 53 insertions(+), 25 deletions(-)
 create mode 100644 client/starter.vbs

diff --git a/client/README.md b/client/README.md
index 246365e..47c1cf9 100644
--- a/client/README.md
+++ b/client/README.md
@@ -2,12 +2,17 @@
 ## 描述
 该客户端用于支持 `client` 消息推送方式。
 
-注意，为防止未授权的 WebSocket 连接，使用该消息推送方式你应设置 ACCESS_TOKEN。
+注意，为防止未授权的 WebSocket 连接，推荐设置 ACCESS_TOKEN。
+
+启动参数：`./client.exe --url http://your.domain.com:port/prefix --token private`
+
+也可以通过设置环境变量来传递配置：`MESSAGE_PUSHER_URL` & `MESSAGE_PUSHER_TOKEN`
 
 ## 原理
 客户端启动后与 message-pusher 服务器建立一个 WebSocket 连接，通过该连接接受要推送消息，并调用系统的消息通知接口进行消息通知。
 
 ## 安装
-推荐将其注册为开机启动的服务，二进制包请前往 [Release 页面](https://github.com/songquanpeng/message-pusher/releases)下载。
-
-启动参数：`./client.exe -server your.domain.com:port -prefix admin -token private`
+1. 首先前往 [Release 页面](https://github.com/songquanpeng/message-pusher/releases)下载可执行文件和 vbs 脚本文件（该脚本文件目的是为了隐藏窗口）。
+2. 编辑 vbs 脚本文件，将其中的可执行文件的路径和后面的参数改成符合你的情况的值。
+3. 之后对 vbs 脚本文件右键创建快捷方式。
+5. 之后把该快捷方式放到开机启动文件夹内：`C:\ProgramData\Microsoft\Windows\Start Menu\Programs\StartUp`
diff --git a/client/main.go b/client/main.go
index 71462ec..78b9e81 100644
--- a/client/main.go
+++ b/client/main.go
@@ -2,16 +2,15 @@ package main
 
 import (
 	"flag"
-	"fmt"
 	"github.com/gorilla/websocket"
 	"log"
 	"net/url"
+	"os"
 )
 
 var (
-	server = flag.String("server", "", "message pusher address")
-	prefix = flag.String("prefix", "", "your prefix")
-	token  = flag.String("token", "", "the access token")
+	conn  = flag.String("url", "", "connection url")
+	token = flag.String("token", "", "the access token")
 )
 
 type Verification struct {
@@ -21,29 +20,50 @@ type Verification struct {
 
 func main() {
 	flag.Parse()
-	u := url.URL{Scheme: "ws", Host: *server, Path: "/"}
-	fmt.Printf("Connecting to %s\n", u.String())
-	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+	connString := *conn
+	if connString == "" {
+		connString = os.Getenv("MESSAGE_PUSHER_URL")
+	}
+	if *token == "" {
+		*token = os.Getenv("MESSAGE_PUSHER_TOKEN")
+	}
+	connUrl, err := url.Parse(connString)
 	if err != nil {
-		log.Fatal("Failed to connect server:", err)
+		log.Fatal("Failed to parse connection url", err)
 		return
 	}
-	defer c.Close()
-	verification := &Verification{
-		Prefix: *prefix,
-		Token:  *token,
+	scheme := "ws"
+	if connUrl.Scheme == "https" {
+		scheme = "wss"
 	}
-	_ = c.WriteJSON(verification)
-
+	u := url.URL{Scheme: scheme, Host: connUrl.Host, Path: "/"}
 	for {
-		var message = new(Message)
-		err = c.ReadJSON(message)
+		log.Printf("Connecting to %s...\n", u.String())
+		c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
 		if err != nil {
-			log.Println("Error occurred when read message:", err)
-		} else {
-			log.Println("Get new message")
-			Notify(message)
+			log.Fatal("Failed to connect to server:", err)
+			return
+		}
+		log.Printf("Server connected.\n")
+		verification := &Verification{
+			Prefix: connUrl.Path[1:],
+			Token:  *token,
+		}
+		err = c.WriteJSON(verification)
+		if err != nil {
+			log.Fatal(err.Error())
+		}
+		for {
+			var message = new(Message)
+			err = c.ReadJSON(message)
+			if err != nil {
+				log.Println("Error occurred when read message:", err)
+				log.Println("Connection lost, retrying...")
+				break
+			} else {
+				log.Println("New message arrived.")
+				Notify(message)
+			}
 		}
 	}
-
 }
diff --git a/client/starter.vbs b/client/starter.vbs
new file mode 100644
index 0000000..1ef3023
--- /dev/null
+++ b/client/starter.vbs
@@ -0,0 +1,3 @@
+Set WshShell = CreateObject("WScript.Shell")
+WshShell.Run chr(34) & "C:\Users\song\Projects\message-pusher\client\message-receiver.exe ""--url https://server_host/prefix" & Chr(34), 0
+Set WshShell = Nothing
\ No newline at end of file