|
|
@@ -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.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.Println("Error occurred when read message:", err)
|
|
|
- } else {
|
|
|
- log.Println("Get new message")
|
|
|
- Notify(message)
|
|
|
+ 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)
|
|
|
+ }
|
|
|
}
|
|
|
}
|
|
|
-
|
|
|
}
|