Merge pull request #388 from daimoc/develop

Add websocket_origin parameters to customize Origin header
processone · Feb 26, 2023 · d2f5135 · d2f5135
2 parents 19c3322 + da21e33
commit d2f5135
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 7 deletions.
diff --git a/include/ts_profile.hrl b/include/ts_profile.hrl
@@ -62,6 +62,7 @@
          ip_transparent = false,  % set IP_TRANSPARENT option on the socket
          websocket_path = "/chat",  % for websocket only
          websocket_frame = "binary",  % for websocket only
+         websocket_origin = "",  % for websocket only
          websocket_subprotocols = [],     % for websocket only
          retry_timeout = 10,        % retry sending in milliseconds
          max_retries = 3,           % maximum number of retries
@@ -96,6 +97,7 @@
          retries=0,   % number of connect retries
          hibernate = 10000, % hibernate if thinktime is >= to this (10sec by default)
          host,        % hostname (or IP) of remote server
+         origin,      % Origin Header
          port,        % server port
          protocol,    % gen_udp, gen_tcp or ssl
          proto_opts = #proto_opts{},  %

diff --git a/src/tsung/ts_server_websocket_ssl.erl b/src/tsung/ts_server_websocket_ssl.erl
@@ -35,14 +35,14 @@
 -include("ts_config.hrl").
 -include("ts_websocket.hrl").
 
--record(state, {parent, socket = none, accept, host, port, path, opts, version,
+-record(state, {parent, socket = none, accept, host, port, path, opts, version, origin,
                 frame, buffer = <<>>, state = not_connected, subprotos = []}).
 
--record(ws_config, {path, version = "13", frame, subprotos}).
+-record(ws_config, {path, version = "13", frame, origin, subprotos}).
 
 protocol_options(#proto_opts{tcp_rcv_size = Rcv, tcp_snd_size = Snd,
-                             websocket_path = Path, websocket_frame = Frame, websocket_subprotocols = SubProtocols}) ->
-    [#ws_config{path = Path, frame = Frame, subprotos = SubProtocols},
+                             websocket_path = Path, websocket_frame = Frame, websocket_origin = Origin, websocket_subprotocols = SubProtocols}) ->
+    [#ws_config{path = Path, frame = Frame, origin = Origin, subprotos = SubProtocols},
      binary,
      {active, once},
      {recbuf, Rcv},
@@ -58,12 +58,13 @@ connect(Host, Port, Opts, Timeout) ->
     Version = WSConfig#ws_config.version,
     Frame = WSConfig#ws_config.frame,
     Protocol = WSConfig#ws_config.subprotos,
+    Origin = WSConfig#ws_config.origin,
 
     case ssl:connect(Host, Port, opts_to_tcp_opts(TcpOpts),Timeout) of
         {ok, Socket} ->
             Pid = spawn_link(
                     fun() ->
-                            loop(#state{parent = Parent, host = Host, port = Port, subprotos = Protocol,
+                            loop(#state{parent = Parent, host = Host, port = Port, subprotos = Protocol, origin = Origin,
                                         opts = TcpOpts, path = Path, version = Version,
                                         frame = Frame, socket = Socket})
                     end),
@@ -75,9 +76,8 @@ connect(Host, Port, Opts, Timeout) ->
     end.
 
 loop(#state{socket = Socket, host = Host, path = Path,
-            version = Version, subprotos = SubProtocol,
+            version = Version, subprotos = SubProtocol, origin = Origin,
             state = not_connected} = State)->
-    Origin = "",
     Headers = "",
     {Handshake, Accept} = websocket:get_handshake(Host, Path, SubProtocol, Version, Origin, Headers),
     ssl:send(Socket, Handshake),

diff --git a/src/tsung_controller/ts_config.erl b/src/tsung_controller/ts_config.erl
@@ -950,6 +950,12 @@ parse(Element = #xmlElement{name=option, attributes=Attrs},
                     NewProto =  OldProto#proto_opts{websocket_subprotocols=SubProtocols},
                     lists:foldl( fun parse/2, Conf#config{proto_opts=NewProto},
                                  Element#xmlElement.content);
+               "websocket_origin" ->
+                   Origin = getAttr(string,Attrs, value, ?config(websocket_origin)),
+                   OldProto =  Conf#config.proto_opts,
+                   NewProto =  OldProto#proto_opts{websocket_origin=Origin},
+                   lists:foldl( fun parse/2, Conf#config{proto_opts=NewProto},
+                                Element#xmlElement.content);
                 "bosh_path" ->
                     Path = getAttr(string,Attrs, value, ?config(bosh_path)),
                     OldProto =  Conf#config.proto_opts,