yaws_revproxy.erl

%%%-------------------------------------------------------------------
%%% File    : yaws_revproxy.erl
%%% Author  :  <klacke@hyber.org>
%%% Description : reverse proxy
%%%
%%% Created :  3 Dec 2003 by  <klacke@hyber.org>
%%%-------------------------------------------------------------------
-module(yaws_revproxy).
-compile(export_all).

-include("../include/yaws.hrl").
-include("../include/yaws_api.hrl").
-include("yaws_debug.hrl").

%% reverse proxy implementation.

%% TODO: Activate proxy keep-alive with a new option ?
-define(proxy_keepalive, false).

%% Proxy socket definition
-record(psock, {s,      %% the socket
		mode,   %% chunked,len,undefined,expectheaders,expectchunked
		prefix, %% The prefix to strip and add
		url,    %% The url we're proxying to
		type,   %% client | server
		r_req,  %% if'we're server, what req method are we processing
		r_host, %% and value of Host: for the cli request
		httpconnection="keep-alive", %% Do we need to keep the connection open ("keep-alive" | "close")  
		state}).%% various depending on mode, ......

%% MREMOND: TODO: Check if redirection works properly (this is
%% fundamental as most of the time cookie based authentication works
%% with redirection).
%% I did not yet check that.

init(CliSock, ARG, _DecPath, _QueryPart, {Prefix, URL}, N) ->
    GC=get(gc), SC=get(sc),
    %% Connect to the backend server
    case connect_url(URL, SC) of
	{ok, Ssock} ->
	    Headers0 = ARG#arg.headers,
            KeepAlive = Headers0#headers.connection,
	    ?Debug("Keep-alive: ~p~n",[KeepAlive]),
	    Cli = sockmode(Headers0, ARG#arg.req,
			   #psock{s = CliSock, prefix = Prefix,
				  url = URL, type = client}),
	    ?Debug("CLI: ~s~n",[?format_record(Cli, psock)]),
	    Headers = rewrite_headers(Cli, Headers0),
	    ReqStr = yaws_api:reformat_request(
		       rewrite_path(ARG#arg.req, Prefix)),
	    Hstr = headers_to_str(Headers),

	    %% Sending received data from the client to the proxied server:
	    %%  MREMOND: TODO: Refactor.
	    yaws:gen_tcp_send(Ssock, [ReqStr, "\r\n", Hstr, "\r\n"]),
	    RemainingData = if
		       N /= 0 ->
			   %% Send data to the server:
			   ClientData=ARG#arg.clidata,
			   yaws:gen_tcp_send(Ssock,ClientData),
				    N - size(ClientData);
		       true ->
			   %% MREMOND: Check this (If N == 0, it might
			   %% be because, we do not know the size of
			   %% the data to receive. For now we will get
			   %% back in http mode, expecting headers
			   %% (new request)
			   0
		   end,

	    Cli2 = case RemainingData of
		       0      -> 
			   %% Sockmode will be changed in the ploop function
			   Cli#psock{mode = expectheaders, state = undefined};
		       _Other -> Cli
		   end,
	    
	    Srv = #psock{s = Ssock, 
			 prefix = Prefix, 
			 url = URL,
			 r_req = (ARG#arg.req)#http_request.method,
			 r_host = (ARG#arg.headers)#headers.host,
			 mode = expectheaders, type = server,
			 httpconnection= if KeepAlive == undefined ->
						 "keep-alive";
					    true ->
						 yaws:lowercase(KeepAlive)
					 end
			},     
	    %% Need to check if we could close the connection after server answer
	    %% Now we _must_ spawn a process here, because we
	    %% can't use {active, once} due to the inefficencies
	    %% that would occur with chunked encodings
	    P1 = proc_lib:spawn_link(?MODULE, ploop, [Cli2, Srv, GC, SC, self()]),
	    ?Debug("Client=~p, Srv=~p", [P1, self()]),
	    ploop(Srv, Cli2, GC, SC, P1);
	_ERR ->
	    yaws:outh_set_dyn_headers(ARG#arg.req, ARG#arg.headers,
				      #urltype{}),
	    yaws_server:deliver_dyn_part(
	      CliSock,  
	      0, "404",
	      0,
	      ARG, no_UT_defined,
	      fun(A)->(SC#sconf.errormod_404):out404(A,get(gc),get(sc)) 
	      end,
	      fun(A)->yaws_server:finish_up_dyn_file(A, CliSock)
	      end
	     )
    end.



headers_to_str(Headers) ->
    lists:map(
      fun(H) -> [H, "\r\n"] end,
      yaws_api:reformat_header(Headers)).


connect_url(URL, _SC) ->
    Port = if
	       URL#url.port == undefined -> 80;
	       true -> URL#url.port
	   end,
    %% FIXME, do ssl:connect here if is_ssl(SC)
    %% as it is now, SSL doesn't work at all here
    gen_tcp:connect(URL#url.host, Port, [{active, false},
					 {packet, http}]).



rewrite_path(Req, Pref) ->
    {abs_path, P} = Req#http_request.path,
    New = Req#http_request{path = {abs_path,strip_prefix(P, Pref)}},
    ?Debug("Rewritten path=~p~nP=~p Pref=~p", [New,P,Pref]),
    New.


strip_prefix(P,"/") ->
    P;
strip_prefix([H|T1],[H|T2]) ->
    strip_prefix(T1,T2).



%% Once we have read the headers, what comes after
%% the headers, 
%% This is applicable both for cli and srv sockets

sockmode(H,Req,Psock) ->
    case Psock#psock.type of
	server ->
	    s_sockmode(H, Req, Psock);
	client ->
	    cont_len_check(H,Psock)
    end.

s_sockmode(H,Resp,Psock) ->
    if Psock#psock.r_req == 'HEAD' ->
	    %% we're replying to a HEAD
	    %% no body
	    Psock#psock{mode = expectheaders, state = undefined};
	
	true ->
	    case lists:member(Resp#http_response.status, 
			      [100,204,205,304,406]) of
		true ->
		    %% no body, illegal
		    Psock#psock{mode = expectheaders, state = undefined};
		false ->
		    cont_len_check(H, Psock)
	    end
    end.


cont_len_check(H,Psock) ->
    case H#headers.content_length of
	undefined ->
	    case H#headers.transfer_encoding of
		"chunked" ->
		    Psock#psock{mode = expectchunked, state = init};
		undefined when Psock#psock.type == client ->
		    Psock#psock{mode = expectheaders, state = undefined};
		undefined when Psock#psock.type == server ->
		    Psock#psock{mode = undefined, state = undefined}
	    end;
	Int when integer(Int) ->
	    Psock#psock{mode = len,
			state = Int};
	List when list(List) ->
	    Psock#psock{mode = len,
			state = list_to_integer(List)}
    
    end.



set_sock_mode(PS) ->
    S = PS#psock.s,
    case PS#psock.mode of
	expectheaders ->
	    inet:setopts(S, [{packet, http}]);
	expectchunked ->
	    inet:setopts(S, [binary, {packet, line}]);
	expect_nl_before_chunked ->
	    inet:setopts(S, [binary, {packet, line}]);
	_ ->
	    inet:setopts(S, [binary, {packet, raw}])
    end.


get_chunk_num(Fd) ->
    get_chunk_num(Fd,nossl).
get_chunk_num(Fd,SSL) ->
    case yaws:do_recv(Fd, 0, SSL) of
	{ok, Line} ->
	    ?Debug("Get chunk num from line ~p~n",[Line]),
	    erlang:list_to_integer(nonl(Line),16);
	{error, _Rsn} ->
	    exit(normal)
    end.

nonl(B) when binary(B) ->
    nonl(binary_to_list(B));
nonl([10|T]) ->
    nonl(T);
nonl([13|T]) ->
    nonl(T);
nonl([H|T]) ->
    [H|nonl(T)];
nonl([]) ->
    [].
	    


get_chunk(_Fd, N, N, _) ->
    [];
get_chunk(Fd, N, Asz,SSL) ->
    case yaws:do_recv(Fd, N, SSL) of
	{ok, Bin} ->
	    SZ = size(Bin),
	    [Bin|get_chunk(Fd, N, SZ+Asz,SSL)];
	_ ->
	    exit(normal)
    end.



ploop(From0, To, GC, SC, Pid) ->
    put(sc, SC),
    put(gc, GC),
    put(ssl, case SC#sconf.ssl of
		 undefined -> nossl;
		 true -> ssl
	     end),
    ploop(From0, To, Pid).


ploop(From0, To, Pid) ->
    From = receive
	       {cli2srv, Method, Host} ->
		   From0#psock{r_req = Method,
			       r_host = Host}
	   after 0 ->
		   From0
	   end,
    set_sock_mode(From),
    TS = To#psock.s,
    case From#psock.mode of
	expectheaders ->
	    case yaws:http_get_headers(From#psock.s, get(ssl)) of
		{R, H0} ->
		    ?Debug("R = ~p~n",[R]),
		    RStr = 
			if
			    %% FIXME handle bad_request here
			    record(R, http_response) ->
				yaws_api:reformat_response(R);
			    record(R, http_request) ->
				Pid ! {cli2srv, R#http_request.method, 
				       H0#headers.host},
				yaws_api:reformat_request(
				  rewrite_path(R, From#psock.prefix))
			end,
		    Hstr = headers_to_str(H = rewrite_headers(From, H0)),
		    yaws:gen_tcp_send(TS, [RStr, "\r\n", Hstr]),
		    From2 = sockmode(H, R, From),
		    if
			From2#psock.mode == expectchunked ->
			    true;
			true ->
			    yaws:gen_tcp_send(TS,"\r\n")
		    end,
		    ploop(From2, To, Pid);
		closed ->
		    done
	    end;
	expectchunked ->
	    %% read the chunk number, we're in line mode
	    N = get_chunk_num(From#psock.s, get(ssl)),
	    if N == 0 ->
		    ok=eat_crnl(From#psock.s, get(ssl)),
		    yaws:gen_tcp_send(TS,["\r\n0\r\n\r\n"]),
		    ?Debug("SEND final 0 ",[]),
		    ploop_keepalive(From#psock{mode = expectheaders, 
				     state = undefined},To, Pid);
	       true ->
		    ploop(From#psock{mode = chunk, 
				    state = N},To, Pid)
	    end;
	chunk ->
	    CG = get_chunk(From#psock.s,From#psock.state, 0, get(ssl)),
	    SZ = From#psock.state,
	    Data2 = ["\r\n", yaws:integer_to_hex(SZ),"\r\n", CG],
	    yaws:gen_tcp_send(TS, Data2),
	    ploop(From#psock{mode = expect_nl_before_chunked,
			     state = undefined}, To, Pid);
	expect_nl_before_chunked ->
	    case yaws:do_recv(From#psock.s, 0, get(ssl)) of
		{ok, <<13,10>>} ->
		    yaws:gen_tcp_send(TS,<<13,10>>),
		    ploop(From#psock{mode = expectchunked,
				     state = undefined}, To, Pid);
		_Other ->
		    exit(normal)
	    end;
	len when From#psock.state == 0 ->
	    ploop_keepalive(From#psock{mode = expectheaders,
			     state = undefined},To, Pid);
	len ->
	    case yaws:do_recv(From#psock.s, From#psock.state, get(ssl)) of
		{ok, Bin} ->
		    SZ = size(Bin),
                    ?Debug("Read ~p bytes~n", [SZ]),
		    yaws:gen_tcp_send(TS, Bin),
		    ploop(From#psock{state = From#psock.state - SZ},
			  To, Pid);
		_Rsn ->
                    ?Debug("Failed to read :~p~n", [_Rsn]),  
		    exit(normal)
	    end;
	undefined ->
	    case yaws:do_recv(From#psock.s, From#psock.state, get(ssl)) of
		{ok, Bin} ->
		    yaws:gen_tcp_send(TS, Bin),
		    ploop(From, To, Pid);
		_ ->
		    exit(normal)
	    end
    end.

%% Before reentering the ploop in expect_header mode (new request/reply),
%% We must check the if we need to keep the connection alive
%% or if we must close it.
ploop_keepalive(_From = #psock{httpconnection="close"}, _To, _Pid) ->
    ?Debug("Connection closed by proxy: No keep-alive~n",[]),
    done;    %%  Close the connection
ploop_keepalive(From, To, Pid) ->
    %% Check server reverse proxy keep-alive setting
    %% TODO: We should get this value from the config file
    case check_server_keepalive() of
	false -> done; %% Close the connection: Server config do not
                       %%  allow proxy keep-alive
	true -> ploop(From, To, Pid) %% Try keeping the connection
                                     %% alive => Wait for headers
    end.

%% TODO: Get proxy keepalive value in SC record
check_server_keepalive() ->
    ?proxy_keepalive.

%% On the way from the client to the server, we need
%% rewrite the Host header

rewrite_headers(PS, H) when PS#psock.type == client ->
    Host =
	if H#headers.host == undefined ->
		undefined;
	   true ->
		U = PS#psock.url,
		[U#url.host,
		 if
		     U#url.port == undefined ->
			 [];
		     true ->
			 [$: | integer_to_list(U#url.port)]
		 end]
	end,

    ?Debug("New Host hdr: ~p~n", [Host]),
    H#headers{host = Host};



%% And on the way from the server to the client we 
%% need to rewrite the Location header, and the
%% Set-Cookie header

rewrite_headers(PS, H) when PS#psock.type == server ->
    ?Debug("Location header to rewrite:  ~p~n", [H#headers.location]),
    Loc = if
	      H#headers.location == undefined ->
		  undefined;
	      true ->
		  ?Debug("parse_url(~p)~n", [H#headers.location]),
		  LocUrl = (catch yaws_api:parse_url(H#headers.location)),
		  ProxyUrl = PS#psock.url,
		  if
		      LocUrl#url.host == ProxyUrl#url.host,
		      LocUrl#url.port == ProxyUrl#url.port,
		      LocUrl#url.scheme == ProxyUrl#url.scheme ->
			  rewrite_loc_url(LocUrl, PS);
		      
		      element(1, LocUrl) == 'EXIT' ->
			  rewrite_loc_rel(PS, H#headers.location);
		      true ->
			  ?Debug("Not rew ~p~n~p~n", 
			      [LocUrl, ProxyUrl]),
			  H#headers.location
		  end
	  end,
    ?Debug("New loc=~p~n", [Loc]),

    %% And we also should do cookies here ... FIXME

    H#headers{location = Loc}.


%% Rewrite a properly formatted location redir
rewrite_loc_url(LocUrl, PS) ->
    SC=get(sc),
    Scheme = yaws_server:redirect_scheme(SC),
    RedirHost = yaws_server:redirect_host(SC, PS#psock.r_host),
    _RealPath = LocUrl#url.path,
    [Scheme, RedirHost, yaws:slash_append(PS#psock.prefix, LocUrl#url.path)].


%% This is the case for broken webservers that reply with
%% Location: /path
%% or even worse, Location: path

rewrite_loc_rel(PS, Loc) ->
    SC=get(sc),
    Scheme = yaws_server:redirect_scheme(SC),
    RedirHost = yaws_server:redirect_host(SC, PS#psock.r_host),
    [Scheme, RedirHost,Loc].



eat_crnl(Fd) ->
    eat_crnl(Fd,nossl).
eat_crnl(Fd,SSL) ->
    yaws:setopts(Fd, [{packet, line}],SSL),
    case yaws:do_recv(Fd,0, SSL) of
	{ok, <<13,10>>} ->
	    ok;
	{ok, [13,10]} ->
	    ok;
	Err ->
	    {error, Err}
    end.