bagas/tunnel-please
Public Access
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

fix: correct read/write handling in CustomWriter
All checks were successful
Docker Build and Push / build-and-push (push) Successful in 5m7s
Details

Browse Source
This commit is contained in:
bagas
2025-12-06 22:17:55 +07:00
parent 0b8bc1dbba
commit 69c3e78728
7 changed files with 137 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
session/handler.go
View File

@ -160,21 +160,29 @@ func (s *SSHSession) HandleHTTPForward(req *ssh.Request, portToBind uint16) {
err := req.Reply(false, nil)
if err != nil {
log.Println("Failed to reply to request:", err)
return
}
return
}
if !registerClient(slug, s) {
log.Printf("Failed to register client with slug: %s", slug)
err := req.Reply(false, nil)
if err != nil {
log.Println("Failed to reply to request:", err)
}
return
}
s.SlugManager.Set(slug)
registerClient(slug, s)
buf := new(bytes.Buffer)
err := binary.Write(buf, binary.BigEndian, uint32(80))
err := binary.Write(buf, binary.BigEndian, uint32(portToBind))
if err != nil {
log.Println("Failed to reply to request:", err)
log.Println("Failed to write port to buffer:", err)
unregisterClient(slug)
return
}
log.Printf("HTTP forwarding approved on port: %d", 80)
log.Printf("HTTP forwarding approved on port: %d", portToBind)
domain := utils.Getenv("domain")
protocol := "http"
@ -184,9 +192,11 @@ func (s *SSHSession) HandleHTTPForward(req *ssh.Request, portToBind uint16) {
s.Interaction.ShowWelcomeMessage()
s.Interaction.SendMessage(fmt.Sprintf("Forwarding your traffic to %s://%s.%s\r\n", protocol, slug, domain))
err = req.Reply(true, buf.Bytes())
if err != nil {
log.Println("Failed to reply to request:", err)
unregisterClient(slug)
return
}
}
@ -194,7 +204,6 @@ func (s *SSHSession) HandleHTTPForward(req *ssh.Request, portToBind uint16) {
func (s *SSHSession) HandleTCPForward(req *ssh.Request, addr string, portToBind uint16) {
s.Forwarder.SetType(types.TCP)
log.Printf("Requested forwarding on %s:%d", addr, 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))
@ -209,25 +218,36 @@ func (s *SSHSession) HandleTCPForward(req *ssh.Request, addr string, portToBind
}
return
}
s.Forwarder.SetListener(listener)
s.Forwarder.SetForwardedPort(portToBind)
s.Interaction.ShowWelcomeMessage()
s.Interaction.SendMessage(fmt.Sprintf("Forwarding your traffic to tcp://%s:%d \r\n", utils.Getenv("domain"), s.Forwarder.GetForwardedPort()))
go s.Forwarder.AcceptTCPConnections()
buf := new(bytes.Buffer)
err = binary.Write(buf, binary.BigEndian, uint32(portToBind))
if err != nil {
log.Println("Failed to reply to request:", err)
log.Println("Failed to write port to buffer:", err)
err = listener.Close()
if err != nil {
log.Printf("Failed to close listener: %s", err)
return
}
return
}
log.Printf("TCP forwarding approved on port: %d", portToBind)
err = req.Reply(true, buf.Bytes())
if err != nil {
log.Println("Failed to reply to request:", err)
err = listener.Close()
if err != nil {
log.Printf("Failed to close listener: %s", err)
return
}
return
}
s.Forwarder.SetListener(listener)
s.Forwarder.SetForwardedPort(portToBind)
s.Interaction.ShowWelcomeMessage()
s.Interaction.SendMessage(fmt.Sprintf("Forwarding your traffic to tcp://%s:%d \r\n", utils.Getenv("domain"), s.Forwarder.GetForwardedPort()))
go s.Forwarder.AcceptTCPConnections()
}
func generateUniqueSlug() string {

19
session/session.go
View File

@ -30,6 +30,8 @@ type SSHSession struct {
Interaction interaction.Controller
Forwarder forwarder.ForwardingController
SlugManager slug.Manager
channelOnce sync.Once
}
func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan ssh.NewChannel) {
@ -61,7 +63,7 @@ func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan
interactionManager.SetSlugModificator(updateClientSlug)
forwarderManager.SetLifecycle(lifecycleManager)
lifecycleManager.SetUnregisterClient(unregisterClient)
session := &SSHSession{
Lifecycle: lifecycleManager,
Interaction: interactionManager,
@ -73,20 +75,23 @@ func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan
go session.Lifecycle.WaitForRunningStatus()
for channel := range sshChan {
ch, reqs, _ := channel.Accept()
if session.Lifecycle.GetChannel() == nil {
ch, reqs, err := channel.Accept()
if err != nil {
log.Printf("failed to accept channel: %v", err)
continue
}
session.channelOnce.Do(func() {
session.Lifecycle.SetChannel(ch)
session.Interaction.SetChannel(ch)
session.Lifecycle.SetStatus(types.SETUP)
go session.HandleGlobalRequest(forwardingReq)
}
})
go session.HandleGlobalRequest(reqs)
}
err := session.Lifecycle.Close()
if err != nil {
if err := session.Lifecycle.Close(); err != nil {
log.Printf("failed to close session: %v", err)
}
return
}()
}