fix: potential resource leak #38
@ -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
|
||||
}
|
||||
|
||||
@ -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
|
||||
//}
|
||||
|
||||
@ -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) {
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -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" +
|
||||
|
||||
Reference in New Issue
Block a user