7d7780a

分支 (1)

提交

管理

管理

master

7d7780a

rc
/
rcproxy
/
rcproxy-main.go

// Copyright 2015 The Gorilla WebSocket Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +  build ignore

package main

import (
	"encoding/json"
	"flag"
	"gitee.com/libbylg/rc/rcproxy/rcproxy"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"html/template"
	"io/ioutil"
	"log"
	"net"
	"net/http"
)

var addr = flag.String("addr", "localhost:8089", "http service address")

var upgrader = websocket.Upgrader{} // use default options

var globalRCProxy = rcproxy.NewRCProxy()

func echo(w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer c.Close()
	for {
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}
		log.Printf("recv: %s", message)
		err = c.WriteMessage(mt, message)
		if err != nil {
			log.Println("write:", err)
			break
		}
	}
}

func home(w http.ResponseWriter, r *http.Request) {
	homeTemplate.Execute(w, "ws://"+r.Host+"/echo")
}

func shellrWebSocket(w http.ResponseWriter, r *http.Request) {
	shellName := ""
	if r.Method == "GET" {
		params := mux.Vars(r)
		shellName = params["SHELLNAME"]
	}

	log.Printf("Client [%s] connected:", shellName)

	//	先升级为 websocket
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}

	//mt, message, err := c.ReadMessage()
	//if err != nil {
	//	log.Println("read:", err)
	//	break
	//}

	//	如果是首次注册
	globalRCProxy.RegisterShellr(rcproxy.NewShellr(shellName, c))
}

func runnerRemoteCall(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		params := mux.Vars(r)

		shellName := params["SHELLNAME"]

		//	先通过URL中的参数找到 Shellr 对象
		shellr, err := globalRCProxy.Find(shellName)
		if err != nil {
			return
		}

		//	读取所有的消息内容
		text, err := ioutil.ReadAll(r.Body)
		if err != nil {
			return
		}

		str := string(text)

		log.Printf("Call: %s/%s: %s\n", r.Host, r.RequestURI, str)

		//	丢给找到的 shellr 对象处理
		repText, err := shellr.Call(text)
		if err != nil {
			return
		}

		//	返回结果回写给调用端
		w.Write(repText)
	}
}

func shellrList(w http.ResponseWriter, r *http.Request) {
	shellrs := make(map[string]string)
	globalRCProxy.EnumShellr(func(name string, addr net.Addr) int {
		shellrs[name] = addr.String()
		return 0
	})

	w.Header()

	text, _ := json.Marshal(shellrs)
	w.Write(text)
}

func main() {
	flag.Parse()
	log.SetFlags(0)

	//	Shellr 侧连接入口（Shellr 需要使用 WebSocket）
	router := mux.NewRouter()
	router.HandleFunc("/shell/{SHELLNAME:[A-Za-z0-9.@_-]+}", shellrWebSocket)
	router.HandleFunc("/shell", shellrList).Methods("GET")

	//	Runner 侧请求入口（Runner 侧不使用 WebSocket）
	router.HandleFunc("/runner/{SHELLNAME:[A-Za-z0-9.@_-]+}", runnerRemoteCall).Methods("POST")
	router.HandleFunc("/echo", echo)
	router.HandleFunc("/", home)
	http.Handle("/", router)

	log.Fatal(http.ListenAndServe(*addr, nil))
}

var homeTemplate = template.Must(template.New("").Parse(`
<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<script>
window.addEventListener("load", function(evt) {
    var output = document.getElementById("output");
    var input = document.getElementById("input");
    var ws;
    var print = function(message) {
        var d = document.createElement("div");
        d.textContent = message;
        output.appendChild(d);
        output.scroll(0, output.scrollHeight);
    };
    document.getElementById("open").onclick = function(evt) {
        if (ws) {
            return false;
        }
        addr = document.getElementById("addr").value
		print("ADDRESS:" + addr)
        ws = new WebSocket(addr);
        ws.onopen = function(evt) {
            print("OPEN");
        }
        ws.onclose = function(evt) {
            print("CLOSE");
            ws = null;
        }
        ws.onmessage = function(evt) {
            print("RESPONSE: " + evt.data);
        }
        ws.onerror = function(evt) {
            print("ERROR: " + evt.data);
        }
        return false;
    };
    document.getElementById("send").onclick = function(evt) {
        if (!ws) {
            return false;
        }
        print("SEND: " + input.value);
        ws.send(input.value);
        return false;
    };
    document.getElementById("close").onclick = function(evt) {
        if (!ws) {
            return false;
        }
        ws.close();
        return false;
    };
});
</script>
</head>
<body>
<table>
<tr><td valign="top" width="50%">
<p>Click "Open" to create a connection to the server,
"Send" to send a message to the server and "Close" to close the connection.
You can change the message and send multiple times.
<p>
<form>
<input id="addr" type="text" value="ws://localhost:8089/echo">
<button id="open">Open</button>
<button id="close">Close</button>
<p>
<input id="input" type="text" value="Hello world!">
<button id="send">Send</button>
</form>
</td><td valign="top" width="50%">
<div id="output" style="max-height: 70vh;overflow-y: scroll;"></div>
</td></tr></table>
</body>
</html>
`))