From 3c8e61bc90686096ed77b319cd7ebffda4b4fb54 Mon Sep 17 00:00:00 2001
From: LowK <davidkevinit@gmail.com>
Date: Sun, 21 Aug 2022 22:49:27 +0700
Subject: [PATCH] Added OpExit

---
 doc/server_plugin.md         | 23 +++++++++++++++++++++++
 pkg/msg/msg.go               | 13 +++++++++++++
 pkg/plugin/server/manager.go | 31 +++++++++++++++++++++++++++++++
 pkg/plugin/server/plugin.go  |  3 ++-
 pkg/plugin/server/types.go   |  5 +++++
 server/control.go            | 19 +++++++++++++++++++
 6 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/doc/server_plugin.md b/doc/server_plugin.md
index d73d2439..8fbf9484 100644
--- a/doc/server_plugin.md
+++ b/doc/server_plugin.md
@@ -94,6 +94,29 @@ Client login operation
 }
 ```
 
+#### Exit
+
+Client exit operation. When client cannot connect to server that will notify to admin.
+Can use a new version (OpExit supported at: https://github.com/lowk3v/frp-notify/releases/tag/OpExitSupported)
+
+```
+{
+    "content": {
+        "version": <string>,
+        "hostname": <string>,
+        "os": <string>,
+        "arch": <string>,
+        "user": <string>,
+        "timestamp": <int64>,
+        "privilege_key": <string>,
+        "run_id": <string>,
+        "pool_count": <int>,
+        "metas": map<string>string,
+        "client_address": <string>
+    }
+}
+```
+
 #### NewProxy
 
 Create new proxy
diff --git a/pkg/msg/msg.go b/pkg/msg/msg.go
index ba6dd63e..e5947385 100644
--- a/pkg/msg/msg.go
+++ b/pkg/msg/msg.go
@@ -83,6 +83,19 @@ type LoginResp struct {
 	Error         string `json:"error,omitempty"`
 }
 
+// When frpc exit success, send this message to exit plugin server
+type Exit struct {
+	Version      string            `json:"version,omitempty"`
+	Hostname     string            `json:"hostname,omitempty"`
+	Os           string            `json:"os,omitempty"`
+	Arch         string            `json:"arch,omitempty"`
+	User         string            `json:"user,omitempty"`
+	PrivilegeKey string            `json:"privilege_key,omitempty"`
+	Timestamp    int64             `json:"timestamp,omitempty"`
+	RunID        string            `json:"run_id,omitempty"`
+	Metas        map[string]string `json:"metas,omitempty"`
+}
+
 // When frpc login success, send this message to frps for running a new proxy.
 type NewProxy struct {
 	ProxyName      string            `json:"proxy_name,omitempty"`
diff --git a/pkg/plugin/server/manager.go b/pkg/plugin/server/manager.go
index 47d11d1c..98b3c1be 100644
--- a/pkg/plugin/server/manager.go
+++ b/pkg/plugin/server/manager.go
@@ -26,6 +26,7 @@ import (
 
 type Manager struct {
 	loginPlugins       []Plugin
+	exitPlugins        []Plugin
 	newProxyPlugins    []Plugin
 	closeProxyPlugins  []Plugin
 	pingPlugins        []Plugin
@@ -36,6 +37,7 @@ type Manager struct {
 func NewManager() *Manager {
 	return &Manager{
 		loginPlugins:       make([]Plugin, 0),
+		exitPlugins:        make([]Plugin, 0),
 		newProxyPlugins:    make([]Plugin, 0),
 		closeProxyPlugins:  make([]Plugin, 0),
 		pingPlugins:        make([]Plugin, 0),
@@ -48,6 +50,9 @@ func (m *Manager) Register(p Plugin) {
 	if p.IsSupport(OpLogin) {
 		m.loginPlugins = append(m.loginPlugins, p)
 	}
+	if p.IsSupport(OpExit) {
+		m.exitPlugins = append(m.exitPlugins, p)
+	}
 	if p.IsSupport(OpNewProxy) {
 		m.newProxyPlugins = append(m.newProxyPlugins, p)
 	}
@@ -99,6 +104,32 @@ func (m *Manager) Login(content *LoginContent) (*LoginContent, error) {
 	return content, nil
 }
 
+func (m *Manager) Exit(content *ExitContent) error {
+	if len(m.exitPlugins) == 0 {
+		return nil
+	}
+
+	errs := make([]string, 0)
+	reqid, _ := util.RandID()
+	xl := xlog.New().AppendPrefix("reqid: " + reqid)
+	ctx := xlog.NewContext(context.Background(), xl)
+	ctx = NewReqidContext(ctx, reqid)
+
+	for _, p := range m.exitPlugins {
+		_, _, err := p.Handle(ctx, OpExit, *content)
+		if err != nil {
+			xl.Warn("send Exit request to plugin [%s] error: %v", p.Name(), err)
+			errs = append(errs, fmt.Sprintf("[%s]: %v", p.Name(), err))
+		}
+	}
+
+	if len(errs) > 0 {
+		return fmt.Errorf("send Exit request to plugin errors: %s", strings.Join(errs, "; "))
+	} else {
+		return nil
+	}
+}
+
 func (m *Manager) NewProxy(content *NewProxyContent) (*NewProxyContent, error) {
 	if len(m.newProxyPlugins) == 0 {
 		return content, nil
diff --git a/pkg/plugin/server/plugin.go b/pkg/plugin/server/plugin.go
index 0d34de54..5d337eab 100644
--- a/pkg/plugin/server/plugin.go
+++ b/pkg/plugin/server/plugin.go
@@ -19,9 +19,10 @@ import (
 )
 
 const (
-	APIVersion = "0.1.0"
+	APIVersion = "0.1.1"
 
 	OpLogin       = "Login"
+	OpExit        = "Exit"
 	OpNewProxy    = "NewProxy"
 	OpCloseProxy  = "CloseProxy"
 	OpPing        = "Ping"
diff --git a/pkg/plugin/server/types.go b/pkg/plugin/server/types.go
index d7d98cb6..90e7bfd4 100644
--- a/pkg/plugin/server/types.go
+++ b/pkg/plugin/server/types.go
@@ -37,6 +37,11 @@ type LoginContent struct {
 	ClientAddress string `json:"client_address,omitempty"`
 }
 
+type ExitContent struct {
+	ClientAddress string `json:"client_address,omitempty"`
+	msg.Exit
+}
+
 type UserInfo struct {
 	User  string            `json:"user"`
 	Metas map[string]string `json:"metas"`
diff --git a/server/control.go b/server/control.go
index 197c94de..829d0377 100644
--- a/server/control.go
+++ b/server/control.go
@@ -394,6 +394,25 @@ func (ctl *Control) stoper() {
 
 	ctl.allShutdown.Done()
 	xl.Info("client exit success")
+
+	notifyContent := &plugin.ExitContent{
+		ClientAddress: ctl.conn.RemoteAddr().String(),
+		Exit: msg.Exit{
+			Version:      ctl.loginMsg.Version,
+			Hostname:     ctl.loginMsg.Hostname,
+			Os:           ctl.loginMsg.Os,
+			Arch:         ctl.loginMsg.Arch,
+			User:         ctl.loginMsg.User,
+			PrivilegeKey: ctl.loginMsg.PrivilegeKey,
+			Timestamp:    ctl.loginMsg.Timestamp,
+			RunID:        ctl.loginMsg.RunID,
+			Metas:        ctl.loginMsg.Metas,
+		},
+	}
+	go func() {
+		_ = ctl.pluginManager.Exit(notifyContent)
+	}()
+
 	metrics.Server.CloseClient()
 }