#!/usr/bin/ruby
# vim: noai:ts=4:sw=4
#
#  Copyright (C) 2010-2017 by Luiz Angelo Daros de Luca
#    luizluca@gmail.com
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    Created on 2010-08-10
#
# Changelog:
#
# * 2017-03-17
# - Add support for http proxy authentication using auto-detect support
#
# * 2010-09-01
# - Add TCP keepalive flags for local socket (thanks Jan Adámek)
# - Increased local buffer to 640k
# - Fixed sync problem with stdin/out when using as ProxyCommand inside ssh config
# - Reconnection on client-webrick disconnection
# - Fixed proxy absence
# - Specified webrick timeout (10m)
# - Treated more signals
# - Divide connection (open, read, write)
#
#
#################
#
# Port Forward client/server over HTTP
#
#################

class String
    def strip_heredoc
        min = scan(/^[ \t]*(?=\S)/).min
        if min == nil
            indent=0
        else
            indent=min.size
        end
        gsub(/^[ \t]{#{indent}}/, '')
    end
end

case ARGV.size
    when 2
        server=true
        location=ARGV[1]
     when 4
        server=false
        url=ARGV[1]
        remote_host=ARGV[2]
        remote_port=ARGV[3].to_i
    else
    $stderr.puts <<-EOF.strip_heredoc
	Use as server:
	#$0 localPort bridgeLocation
	
	Use as client:
	#$0 localPort|STDIN bridgeURL remoteAddress remotePort
	
	Ex:
	bridge$ #$0 8080 /bridge
	
	client$ #$0 8022 http://mybridge:8080/bridge remoteServer 22
	client$ ssh localhost -p 8022
	EOF
    exit
end

localPort=ARGV[0]

require 'socket'
require 'thread'
Thread.abort_on_exception = true

MAXBUF=1024*640 # 640k
READ_TIMEOUT=60
if server
    require 'webrick'

    class BridgeServlet < WEBrick::HTTPServlet::AbstractServlet
        @@connections = {}

        def get_socket(req,res)
            conn_id = req.path.split("/").last
            conn = @@connections[conn_id]
            if not conn
                $stderr.puts "[#{conn_id}] Connection is not open for #{conn_id}=#{conn}"
                res.status=404
                return nil
            end
            conn
        end

        def close_socket(req,res)
            conn_id = req.path.split("/").last
            return if not conn=@@connections[conn_id]
            conn.close if not conn.closed?
            @@connections.delete(conn_id)
        end

        def do_POST(req, res)
            conn_id = req.path.split("/").last
            (remote_host,remote_port)=req.body.split(":")
            remote_port=remote_port.to_i
            $stderr.puts "[#{conn_id}] Opening connection to #{remote_host}:#{remote_port} for #{req.peeraddr}..."
            begin
                conn=TCPSocket.open(remote_host,remote_port)
                conn.setsockopt(Socket::SOL_SOCKET, Socket::SO_KEEPALIVE, true)
                @@connections[conn_id]=conn
                $stderr.puts "[#{conn_id}] Connected to #{remote_host}:#{remote_port} #{conn}"
                res.status=201
            rescue
                res['Content-Type'] = "text/plain"
                res.body=$!.message
                $stderr.puts "[#{conn_id}] Connection failed: #$!"
                res.status=406 # Not Acceptable
            end
        end

        def do_PUT(req, res)
            conn_id = req.path.split("/").last
            return if not s=get_socket(req,res)
            begin
                s.print req.body
            rescue Errno::EPIPE, IOError
                $stderr.puts "[#{conn_id}] Connection closed in remote destination (PUT)"
                close_socket(req, res)
                res.status=410 # GONE
            end
        end

        def do_GET(req, res)
            conn_id = req.path.split("/").last
            return if not conn=get_socket(req,res)
            begin
                res.body = conn.read_nonblock(MAXBUF)
            rescue IO::WaitReadable
                if IO.select([conn],[],[],READ_TIMEOUT-3)
                    retry
                else # Timeout
                    res.body = ''
                    res.status=204 # No Content
                end
            rescue EOFError, Errno::EPIPE, IOError
                if conn.closed?
                    $stderr.puts "[#{conn_id}] Connection closed by remote destination (GET)"
                else
                    $stderr.puts "[#{conn_id}] Connection closed (GET) #{$!}"
                end
                close_socket(req, res)
                res.status=410 # Gone
            rescue
                $stderr.puts $!.class
            end
        end

        def do_DELETE(req, res)
            conn_id = req.path.split("/").last
            $stderr.puts "[#{conn_id}] Connection closed in client. Trying to close it in remote destination"
            close_socket(req, res)
            res.status=200
        end

    end

    s = WEBrick::HTTPServer.new(
        :Port            => localPort.to_i,
        :RequestTimeout  => 600
    )
    s.mount(location, BridgeServlet)
    trap("INT"){ s.shutdown }
    s.start

else # Client

    require "net/http"
    require 'uri'

    # HACK! https://bugs.ruby-lang.org/issues/12921
    class Net::HTTP
        def proxy_user
          if @proxy_from_env then
            proxy_uri && proxy_uri.user
          else
            @proxy_user
          end
        end
        def proxy_pass
          if @proxy_from_env then
            proxy_uri && proxy_uri.password
          else
            @proxy_pass
          end
        end
    end

    url = URI(url)
    if url.find_proxy
        $stderr.puts "Using proxy: #{url.find_proxy}"
    end

    case localPort
    when "STDIN","-"
        _in=$stdin
        _out=$stdout
        # Keeps this unbuffered
        _in.sync=true
        _out.sync=true
    else
        $stderr.puts "Opening local port #{localPort}"
        local_server = TCPServer.open(localPort.to_i)
        $stderr.puts "Waiting for local connection to port #{localPort}"
        local_server_conn = local_server.accept
        _in=_out=local_server_conn
        # Keep connection alive
        local_server_conn.setsockopt(Socket::SOL_SOCKET, Socket::SO_KEEPALIVE, true)
    end

    #unique Connection IDentifier
    conn_id=`uuidgen`.chomp
    connected=false


    # Open the connection
    begin
        begin
            $stderr.puts "Opening connection over HTTP bridge for #{remote_host}:#{remote_port} (#{conn_id})"
            # HACK! new().start() because of https://bugs.ruby-lang.org/issues/13351
            Net::HTTP.new(url.host, url.port).start do |http|
                error=nil
                res = http.post("#{url.path}/#{conn_id}","#{remote_host}:#{remote_port}")
                if res.kind_of? Net::HTTPCreated
                    connected=true
                else
                    $stderr.puts "The bridge failed to connect to the remote location (#{res}): #{res.body}"
                    connected=false
                    break
                end
            end
        rescue Errno::EPIPE
            $stderr.puts "Connection to bridge #{url} closed"
            retry
        end

        # Not connected, nothing more to do
        exit 1 if not connected

        # Launch Local write/Remote read loop
        Thread.new do
            $stderr.puts "Local write/Remote read loop started"
            begin
                Net::HTTP.new(url.host, url.port).start do |http|
                    http.read_timeout=READ_TIMEOUT+3
                    while connected
                        res = http.get("#{url.path}/#{conn_id}") {|res| _out.print res }
                        if res.kind_of? Net::HTTPGone
                            $stderr.puts "Connection closed in remote location (lw/rr)"
                            connected=false
                            _in.close if not _in.closed?
                            break
                        elsif res.kind_of? Net::HTTPNotFound
                            $stderr.puts "Connection not opened on bridge"
                            connected=false
                            break
                        end
                    end
                end
            rescue Errno::EPIPE
                _in.close if not _in.closed?
            rescue EOFError
                # retry if local connection is still open
                retry if not _in.closed?
                $stderr.puts "Connection to bridge closed (lw/rr)"
            rescue Errno::ECONNREFUSED
                $stderr.puts "Connection to bridge failed (lw/rr)"
                _in.close if not _in.closed?
            rescue Timeout::Error
                $stderr.puts "Timeout (lw/rr)"
                retry if not _in.closed?
            rescue Net::ReadTimeout
                $stderr.puts "Net::ReadTimeout (lw/rr)"
                retry if not _in.closed?
            rescue Net::HTTPServerError
                raise res.message
            end
        end

        # If CTRL+C, SIGHUP or SIGTERM, close _in (and itself)
        trap("INT"){ _in.close }
        trap("SIGHUP"){ _in.close }
        trap("SIGTERM"){ _in.close }

        # Launch "Local read/Remote write loop started"
        $stderr.puts "Local read/Remote write loop started"
        # Keep buffer between retries
        buf=nil
        begin
            Net::HTTP.new(url.host, url.port).start do |http|
                while connected
                    begin
                        # first read from _in ONLY if buf is empty
                        buf=_in.readpartial(MAXBUF) if !buf
                    rescue IO::WaitReadable
                        IO.select([_in])
                        retry
                    rescue EOFError
                        $stderr.puts "Local connection closed (lr/rw)"
                        $stderr.puts "Closing bridge connection to remote location"
                        http.delete("#{url.path}/#{conn_id}")
                        connected=false
                        break
                    end

                    res = http.put("#{url.path}/#{conn_id}",buf)
                    if res.kind_of? Net::HTTPGone
                        $stderr.puts "Connection closed in remote location (lr/rw)"
                        _in.close if not _in.closed?
                        connected=false
                        break
                    end

                    # Buffer sent, clear it
                    buf=nil
                end
            end
        rescue Net::ReadTimeout
            $stderr.puts "Net::ReadTimeout (lr/rw)"
            retry if not _in.closed? and connected
        rescue Net::HTTPServerError
            raise "Server said: '#{res.message}'"
        rescue Errno::EPIPE, IOError
            # retry if local connection is still open
            retry if not _in.closed?
            $stderr.puts "Connection to bridge closed (lr/rw)"
        end
    rescue Errno::ECONNREFUSED
        $stderr.puts "Connection to bridge failed"
    end
end

$stderr.puts "Program Finished!"
