130 lines
2.7 KiB
Go
130 lines
2.7 KiB
Go
package frpLogger
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
)
|
|
|
|
type TokenValidatorFunc func(token string) (int, error)
|
|
|
|
type LogSocketHandler struct {
|
|
validateToken TokenValidatorFunc
|
|
}
|
|
|
|
func NewLogSocketHandler(validator TokenValidatorFunc) *LogSocketHandler {
|
|
return &LogSocketHandler{
|
|
validateToken: validator,
|
|
}
|
|
}
|
|
|
|
type ErrorResponse struct {
|
|
Success bool `json:"success"`
|
|
Message string `json:"message"`
|
|
}
|
|
|
|
func sendWSError(conn interface {
|
|
WriteJSON(v interface{}) error
|
|
}, message string) {
|
|
resp := ErrorResponse{
|
|
Success: false,
|
|
Message: message,
|
|
}
|
|
conn.WriteJSON(resp)
|
|
}
|
|
|
|
func (h *LogSocketHandler) Handle(w http.ResponseWriter, r *http.Request) {
|
|
instanceID, err := ParseInstanceIDFromQuery(r)
|
|
if err != nil {
|
|
SendHTTPError(w, http.StatusBadRequest, err.Error())
|
|
return
|
|
}
|
|
|
|
token := ParseTokenFromQuery(r)
|
|
if token == "" {
|
|
SendHTTPError(w, http.StatusUnauthorized, "Token is required")
|
|
return
|
|
}
|
|
|
|
userID, err := h.validateToken(token)
|
|
if err != nil {
|
|
SendHTTPError(w, http.StatusUnauthorized, fmt.Sprintf("Invalid token: %v", err))
|
|
return
|
|
}
|
|
|
|
_, err = GetUserByID(userID)
|
|
if err != nil {
|
|
SendHTTPError(w, http.StatusUnauthorized, "User not found")
|
|
return
|
|
}
|
|
|
|
_, err = DBQueryFrpcInstanceByID(instanceID)
|
|
if err != nil {
|
|
SendHTTPError(w, http.StatusNotFound, "Instance not found")
|
|
return
|
|
}
|
|
|
|
// if instance.UserID != userID {
|
|
// SendHTTPError(w, http.StatusForbidden, "You don't have access to this instance")
|
|
// return
|
|
// }
|
|
|
|
serviceName, err := GetServiceNameByInstanceID(instanceID)
|
|
if err != nil {
|
|
SendHTTPError(w, http.StatusInternalServerError, "Failed to get service name")
|
|
return
|
|
}
|
|
|
|
conn, err := upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
log.Printf("[InstanceLogSocketHandler] Failed to upgrade connection: %v", err)
|
|
return
|
|
}
|
|
defer conn.Close()
|
|
|
|
streamer := NewInstanceLogStreamer(instanceID, serviceName, conn)
|
|
streamer.Start()
|
|
|
|
done := make(chan struct{})
|
|
|
|
go func() {
|
|
defer close(done)
|
|
for {
|
|
_, _, err := conn.ReadMessage()
|
|
if err != nil {
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
<-done
|
|
streamer.Stop()
|
|
}
|
|
|
|
func SendHTTPError(w http.ResponseWriter, statusCode int, message string) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(statusCode)
|
|
resp := ErrorResponse{
|
|
Success: false,
|
|
Message: message,
|
|
}
|
|
jsonResp, _ := json.Marshal(resp)
|
|
w.Write(jsonResp)
|
|
}
|
|
|
|
type InstanceLogHandler struct {
|
|
validateToken TokenValidatorFunc
|
|
}
|
|
|
|
func NewInstanceLogHandler(validator TokenValidatorFunc) *InstanceLogHandler {
|
|
return &InstanceLogHandler{
|
|
validateToken: validator,
|
|
}
|
|
}
|
|
|
|
func (h *InstanceLogHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
handler := NewLogSocketHandler(h.validateToken)
|
|
handler.Handle(w, r)
|
|
}
|