diff --git a/main b/main deleted file mode 100644 index 677966f..0000000 Binary files a/main and /dev/null differ diff --git a/server/http.go b/server/http.go index 6ae0454..9eafb7d 100644 --- a/server/http.go +++ b/server/http.go @@ -298,13 +298,22 @@ func forwardRequest(cw *CustomWriter, initialRequest *RequestHeaderFactory, sshS channel, reqs, err := sshSession.Lifecycle.GetConnection().OpenChannel("forwarded-tcpip", payload) if err != nil { log.Printf("Failed to open forwarded-tcpip channel: %v", err) + if closer, ok := cw.writer.(io.Closer); ok { + if closeErr := closer.Close(); closeErr != nil { + log.Printf("Failed to close connection: %v", closeErr) + } + } return } go func() { + defer func() { + if r := recover(); r != nil { + log.Printf("Panic in request handler goroutine: %v", r) + } + }() for req := range reqs { - err := req.Reply(false, nil) - if err != nil { + if err := req.Reply(false, nil); err != nil { log.Printf("Failed to reply to request: %v", err) return } diff --git a/server/middleware.go b/server/middleware.go index a0e3c2b..acbd8bd 100644 --- a/server/middleware.go +++ b/server/middleware.go @@ -44,119 +44,3 @@ func (ff *ForwardedFor) HandleRequest(header *RequestHeaderFactory) error { header.Set("X-Forwarded-For", host) return nil } - -//TODO: Implement caching atau enggak -//const maxCacheSize = 50 * 1024 * 1024 -// -//type DiskCacheMiddleware struct { -// dir string -// mu sync.Mutex -// file *os.File -// path string -// cacheable bool -//} -// -//func NewDiskCacheMiddleware() *DiskCacheMiddleware { -// return &DiskCacheMiddleware{dir: "cache"} -//} -// -//func (c *DiskCacheMiddleware) ensureDir() error { -// return os.MkdirAll(c.dir, 0755) -//} -// -//func (c *DiskCacheMiddleware) cacheKey(method, path string) string { -// return fmt.Sprintf("%s_%s.cache", method, base64.URLEncoding.EncodeToString([]byte(path))) -//} -// -//func (c *DiskCacheMiddleware) filePath(method, path string) string { -// return filepath.Join(c.dir, c.cacheKey(method, path)) -//} -// -//func fileExists(path string) bool { -// _, err := os.Stat(path) -// if err == nil { -// return true -// } -// if os.IsNotExist(err) { -// return false -// } -// return false -//} -// -//func canCacheRequest(header *RequestHeaderFactory) bool { -// if header.Method != "GET" { -// return false -// } -// -// if cacheControl := header.Get("Cache-Control"); cacheControl != "" { -// if strings.Contains(cacheControl, "no-store") || strings.Contains(cacheControl, "private") || strings.Contains(cacheControl, "no-cache") || strings.Contains(cacheControl, "max-age=0") { -// return false -// } -// } -// -// if header.Get("Authorization") != "" { -// return false -// } -// -// if header.Get("Cookie") != "" { -// return false -// } -// -// return true -//} -// -//func (c *DiskCacheMiddleware) HandleRequest(header *RequestHeaderFactory) error { -// if !canCacheRequest(header) { -// c.cacheable = false -// return nil -// } -// -// c.cacheable = true -// _ = c.ensureDir() -// path := c.filePath(header.Method, header.Path) -// -// if fileExists(path + ".finish") { -// c.file = nil -// return nil -// } -// -// if c.file != nil { -// err := c.file.Close() -// if err != nil { -// return err -// } -// err = os.Rename(c.path, c.path+".finish") -// if err != nil { -// return err -// } -// } -// -// c.path = path -// f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) -// if err != nil { -// return err -// } -// -// c.file = f -// -// return nil -//} -// -//func (c *DiskCacheMiddleware) HandleResponse(header *ResponseHeaderFactory, body []byte) error { -// if !c.cacheable { -// return nil -// } -// -// if c.file == nil { -// header.Set("X-Cache", "HIT") -// return nil -// } -// -// _, err := c.file.Write(body) -// if err != nil { -// return err -// } -// -// header.Set("X-Cache", "MISS") -// return nil -//} diff --git a/session/forwarder/forwarder.go b/session/forwarder/forwarder.go index c0330f8..9d94abe 100644 --- a/session/forwarder/forwarder.go +++ b/session/forwarder/forwarder.go @@ -4,7 +4,6 @@ import ( "bytes" "encoding/binary" "errors" - "fmt" "io" "log" "net" @@ -16,12 +15,11 @@ import ( ) type Forwarder struct { - Listener net.Listener - TunnelType types.TunnelType - ForwardedPort uint16 - SlugManager slug.Manager - Lifecycle Lifecycle - ActiveForwarder []chan struct{} + Listener net.Listener + TunnelType types.TunnelType + ForwardedPort uint16 + SlugManager slug.Manager + Lifecycle Lifecycle } type Lifecycle interface { @@ -41,27 +39,6 @@ type ForwardingController interface { SetLifecycle(lifecycle Lifecycle) CreateForwardedTCPIPPayload(origin net.Addr) []byte WriteBadGatewayResponse(dst io.Writer) - AddActiveForwarder(drop chan struct{}) - DropAllForwarder() int - GetForwarderCount() int -} - -func (f *Forwarder) AddActiveForwarder(drop chan struct{}) { - f.ActiveForwarder = append(f.ActiveForwarder, drop) -} - -func (f *Forwarder) DropAllForwarder() int { - total := 0 - for _, d := range f.ActiveForwarder { - close(d) - total += 1 - } - f.ActiveForwarder = nil - return total -} - -func (f *Forwarder) GetForwarderCount() int { - return len(f.ActiveForwarder) } func (f *Forwarder) SetLifecycle(lifecycle Lifecycle) { @@ -82,7 +59,10 @@ func (f *Forwarder) AcceptTCPConnections() { channel, reqs, err := f.Lifecycle.GetConnection().OpenChannel("forwarded-tcpip", payload) if err != nil { log.Printf("Failed to open forwarded-tcpip channel: %v", err) - return + if closeErr := conn.Close(); closeErr != nil { + log.Printf("Failed to close connection: %v", closeErr) + } + continue } go func() { @@ -99,8 +79,7 @@ func (f *Forwarder) AcceptTCPConnections() { } func (f *Forwarder) HandleConnection(dst io.ReadWriter, src ssh.Channel, remoteAddr net.Addr) { - drop := make(chan struct{}) - defer func(src ssh.Channel) { + defer func() { _, err := io.Copy(io.Discard, src) if err != nil { log.Printf("Failed to discard connection: %v", err) @@ -108,34 +87,38 @@ func (f *Forwarder) HandleConnection(dst io.ReadWriter, src ssh.Channel, remoteA err = src.Close() if err != nil && !errors.Is(err, io.EOF) { - log.Printf("Error closing connection: %v", err) + log.Printf("Error closing source channel: %v", err) } - }(src) + + if closer, ok := dst.(io.Closer); ok { + err = closer.Close() + if err != nil && !errors.Is(err, io.EOF) { + log.Printf("Error closing destination connection: %v", err) + } + } + }() + log.Printf("Handling new forwarded connection from %s", remoteAddr) + done := make(chan struct{}, 2) + go func() { _, err := io.Copy(src, dst) if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) { log.Printf("Error copying from conn.Reader to channel: %v", err) } + done <- struct{}{} }() go func() { - select { - case <-drop: - fmt.Println("Closinggggg") - return + _, err := io.Copy(dst, src) + if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) { + log.Printf("Error copying from channel to conn.Writer: %v", err) } + done <- struct{}{} }() - f.AddActiveForwarder(drop) - - _, err := io.Copy(dst, src) - - if err != nil && !errors.Is(err, io.EOF) { - log.Printf("Error copying from channel to conn.Writer: %v", err) - } - return + <-done } func (f *Forwarder) SetType(tunnelType types.TunnelType) { diff --git a/session/handler.go b/session/handler.go index 8e7019e..79ee46c 100644 --- a/session/handler.go +++ b/session/handler.go @@ -211,6 +211,9 @@ func (s *SSHSession) HandleTCPForward(req *ssh.Request, addr string, portToBind listener, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", portToBind)) if err != nil { s.Interaction.SendMessage(fmt.Sprintf("Port %d is already in use or restricted. Please choose a different port.\r\n", portToBind)) + if setErr := portUtil.Manager.SetPortStatus(portToBind, false); setErr != nil { + log.Printf("Failed to reset port status: %v", setErr) + } err = req.Reply(false, nil) if err != nil { log.Println("Failed to reply to request:", err) @@ -227,6 +230,9 @@ func (s *SSHSession) HandleTCPForward(req *ssh.Request, addr string, portToBind err = binary.Write(buf, binary.BigEndian, uint32(portToBind)) if err != nil { log.Println("Failed to write port to buffer:", err) + if setErr := portUtil.Manager.SetPortStatus(portToBind, false); setErr != nil { + log.Printf("Failed to reset port status: %v", setErr) + } err = listener.Close() if err != nil { log.Printf("Failed to close listener: %s", err) @@ -239,6 +245,9 @@ func (s *SSHSession) HandleTCPForward(req *ssh.Request, addr string, portToBind err = req.Reply(true, buf.Bytes()) if err != nil { log.Println("Failed to reply to request:", err) + if setErr := portUtil.Manager.SetPortStatus(portToBind, false); setErr != nil { + log.Printf("Failed to reset port status: %v", setErr) + } err = listener.Close() if err != nil { log.Printf("Failed to close listener: %s", err) diff --git a/session/interaction/interaction.go b/session/interaction/interaction.go index 5800a06..17020a8 100644 --- a/session/interaction/interaction.go +++ b/session/interaction/interaction.go @@ -39,8 +39,6 @@ type Forwarder interface { Close() error GetTunnelType() types.TunnelType GetForwardedPort() uint16 - DropAllForwarder() int - GetForwarderCount() int } type Interaction struct { @@ -118,8 +116,6 @@ func (i *Interaction) handleInteractiveMode(char byte) { switch i.InteractionType { case types.Slug: i.HandleSlugEditMode(char) - case types.Drop: - i.HandleDropMode(char) } } @@ -271,8 +267,17 @@ func (i *Interaction) returnToMainScreen() { } func (i *Interaction) HandleSlugCancel() { + i.SendMessage(clearScreen) + i.SendMessage("\r\n\r\n⚠️ SUBDOMAIN EDIT CANCELLED ⚠️\r\n\r\n") + i.SendMessage("Press any key to continue...\r\n") + i.InteractiveMode = false - i.showMessageAndWait("\r\n\r\n⚠️ SUBDOMAIN EDIT CANCELLED ⚠️\r\n\r\n") + i.InteractionType = "" + i.WaitForKeyPress() + + i.SendMessage(clearScreen) + i.ShowWelcomeMessage() + i.ShowForwardingMessage() } func (i *Interaction) HandleSlugUpdateError() { @@ -295,7 +300,6 @@ func (i *Interaction) HandleCommand(command string) { "/help": i.handleHelpCommand, "/clear": i.handleClearCommand, "/slug": i.handleSlugCommand, - "/drop": i.handleDropCommand, } if handler, exists := handlers[command]; exists { @@ -315,7 +319,7 @@ func (i *Interaction) handleByeCommand() { } func (i *Interaction) handleHelpCommand() { - i.SendMessage("\r\nAvailable commands: /bye, /help, /clear, /slug, /drop\r\n") + i.SendMessage("\r\nAvailable commands: /bye, /help, /clear, /slug\r\n") } func (i *Interaction) handleClearCommand() { @@ -340,13 +344,6 @@ func (i *Interaction) handleSlugCommand() { i.SendMessage("➤ " + i.EditSlug + "." + domain) } -func (i *Interaction) handleDropCommand() { - i.InteractiveMode = true - i.InteractionType = types.Drop - i.SendMessage(clearScreen) - i.ShowDropMessage() -} - func (i *Interaction) ShowForwardingMessage() { domain := utils.Getenv("domain") @@ -361,64 +358,6 @@ func (i *Interaction) ShowForwardingMessage() { } } -func (i *Interaction) HandleDropMode(char byte) { - switch { - case char == enterChar || char == 'y' || char == 'Y': - i.executeDropAll() - case char == escapeChar || char == 'n' || char == 'N' || char == ctrlC: - i.cancelDrop() - } -} - -func (i *Interaction) executeDropAll() { - count := i.Forwarder.DropAllForwarder() - message := fmt.Sprintf("Dropped %d forwarders\r\n", count) - i.showMessageAndWait(message) -} - -func (i *Interaction) cancelDrop() { - i.showMessageAndWait("Dropping canceled.\r\n") -} - -func (i *Interaction) showMessageAndWait(message string) { - i.SendMessage(clearScreen) - i.SendMessage(message) - i.SendMessage("Press any key to continue...\r\n") - - i.InteractiveMode = false - i.InteractionType = "" - i.WaitForKeyPress() - - i.SendMessage(clearScreen) - i.ShowWelcomeMessage() - i.ShowForwardingMessage() -} - -func (i *Interaction) ShowDropMessage() { - confirmText := fmt.Sprintf(" ║ Drop ALL %d active connections?", i.Forwarder.GetForwarderCount()) - boxWidth := calculateBoxWidth(confirmText) - - box := buildDropConfirmationBox(boxWidth, confirmText) - i.SendMessage("\r\n" + box + "\r\n\r\n") -} - -func buildDropConfirmationBox(boxWidth int, confirmText string) string { - topBorder := " ╔" + strings.Repeat("═", boxWidth-4) + "╗\r\n" - title := centerText("DROP CONFIRMATION", boxWidth-4) - header := " ║" + title + "║\r\n" - midBorder := " ╠" + strings.Repeat("═", boxWidth-4) + "╣\r\n" - emptyLine := " ║" + strings.Repeat(" ", boxWidth-4) + "║\r\n" - - confirmLine := confirmText + strings.Repeat(" ", boxWidth-len(confirmText)+1) + "║\r\n" - - controlText := " ║ [Enter/Y] Confirm [N/Esc] Cancel" - controlLine := controlText + strings.Repeat(" ", boxWidth-len(controlText)+1) + "║\r\n" - - bottomBorder := " ╚" + strings.Repeat("═", boxWidth-4) + "╝\r\n" - - return topBorder + header + midBorder + emptyLine + confirmLine + emptyLine + controlLine + emptyLine + bottomBorder -} - func (i *Interaction) ShowWelcomeMessage() { asciiArt := []string{ ` _______ _ _____ _ `, @@ -436,7 +375,6 @@ func (i *Interaction) ShowWelcomeMessage() { ` - '/help' : Show this help message`, ` - '/clear' : Clear the current line`, ` - '/slug' : Set custom subdomain`, - ` - '/drop' : Drop all active forwarders`, } for _, line := range asciiArt { @@ -483,6 +421,10 @@ func (i *Interaction) WaitForKeyPress() { if err == nil { break } + if err != nil { + log.Printf("Error reading keypress: %v", err) + break + } } } diff --git a/types/types.go b/types/types.go index 1ad818c..e0fc74b 100644 --- a/types/types.go +++ b/types/types.go @@ -19,7 +19,6 @@ type InteractionType string const ( Slug InteractionType = "SLUG" - Drop InteractionType = "DROP" ) var BadGatewayResponse = []byte("HTTP/1.1 502 Bad Gateway\r\n" +