1#
2# utils.rb -- Miscellaneous utilities
3#
4# Author: IPR -- Internet Programming with Ruby -- writers
5# Copyright (c) 2001 TAKAHASHI Masayoshi, GOTOU Yuuzou
6# Copyright (c) 2002 Internet Programming with Ruby writers. All rights
7# reserved.
8#
9# $IPR: utils.rb,v 1.10 2003/02/16 22:22:54 gotoyuzo Exp $
10
11require 'socket'
12require 'fcntl'
13begin
14  require 'etc'
15rescue LoadError
16  nil
17end
18
19module WEBrick
20  module Utils
21    ##
22    # Sets IO operations on +io+ to be non-blocking
23    def set_non_blocking(io)
24      flag = File::NONBLOCK
25      if defined?(Fcntl::F_GETFL)
26        flag |= io.fcntl(Fcntl::F_GETFL)
27      end
28      io.fcntl(Fcntl::F_SETFL, flag)
29    end
30    module_function :set_non_blocking
31
32    ##
33    # Sets the close on exec flag for +io+
34    def set_close_on_exec(io)
35      if defined?(Fcntl::FD_CLOEXEC)
36        io.fcntl(Fcntl::F_SETFD, Fcntl::FD_CLOEXEC)
37      end
38    end
39    module_function :set_close_on_exec
40
41    ##
42    # Changes the process's uid and gid to the ones of +user+
43    def su(user)
44      if defined?(Etc)
45        pw = Etc.getpwnam(user)
46        Process::initgroups(user, pw.gid)
47        Process::Sys::setgid(pw.gid)
48        Process::Sys::setuid(pw.uid)
49      else
50        warn("WEBrick::Utils::su doesn't work on this platform")
51      end
52    end
53    module_function :su
54
55    ##
56    # The server hostname
57    def getservername
58      host = Socket::gethostname
59      begin
60        Socket::gethostbyname(host)[0]
61      rescue
62        host
63      end
64    end
65    module_function :getservername
66
67    ##
68    # Creates TCP server sockets bound to +address+:+port+ and returns them.
69    #
70    # It will create IPV4 and IPV6 sockets on all interfaces.
71    def create_listeners(address, port, logger=nil)
72      unless port
73        raise ArgumentError, "must specify port"
74      end
75      res = Socket::getaddrinfo(address, port,
76                                Socket::AF_UNSPEC,   # address family
77                                Socket::SOCK_STREAM, # socket type
78                                0,                   # protocol
79                                Socket::AI_PASSIVE)  # flag
80      last_error = nil
81      sockets = []
82      res.each{|ai|
83        begin
84          logger.debug("TCPServer.new(#{ai[3]}, #{port})") if logger
85          sock = TCPServer.new(ai[3], port)
86          port = sock.addr[1] if port == 0
87          Utils::set_close_on_exec(sock)
88          sockets << sock
89        rescue => ex
90          logger.warn("TCPServer Error: #{ex}") if logger
91          last_error  = ex
92        end
93      }
94      raise last_error if sockets.empty?
95      return sockets
96    end
97    module_function :create_listeners
98
99    ##
100    # Characters used to generate random strings
101    RAND_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
102                 "0123456789" +
103                 "abcdefghijklmnopqrstuvwxyz"
104
105    ##
106    # Generates a random string of length +len+
107    def random_string(len)
108      rand_max = RAND_CHARS.bytesize
109      ret = ""
110      len.times{ ret << RAND_CHARS[rand(rand_max)] }
111      ret
112    end
113    module_function :random_string
114
115    ###########
116
117    require "thread"
118    require "timeout"
119    require "singleton"
120
121    ##
122    # Class used to manage timeout handlers across multiple threads.
123    #
124    # Timeout handlers should be managed by using the class methods which are
125    # synchronized.
126    #
127    #   id = TimeoutHandler.register(10, Timeout::Error)
128    #   begin
129    #     sleep 20
130    #     puts 'foo'
131    #   ensure
132    #     TimeoutHandler.cancel(id)
133    #   end
134    #
135    # will raise Timeout::Error
136    #
137    #   id = TimeoutHandler.register(10, Timeout::Error)
138    #   begin
139    #     sleep 5
140    #     puts 'foo'
141    #   ensure
142    #     TimeoutHandler.cancel(id)
143    #   end
144    #
145    # will print 'foo'
146    #
147    class TimeoutHandler
148      include Singleton
149
150      ##
151      # Mutex used to synchronize access across threads
152      TimeoutMutex = Mutex.new # :nodoc:
153
154      ##
155      # Registers a new timeout handler
156      #
157      # +time+:: Timeout in seconds
158      # +exception+:: Exception to raise when timeout elapsed
159      def TimeoutHandler.register(seconds, exception)
160        TimeoutMutex.synchronize{
161          instance.register(Thread.current, Time.now + seconds, exception)
162        }
163      end
164
165      ##
166      # Cancels the timeout handler +id+
167      def TimeoutHandler.cancel(id)
168        TimeoutMutex.synchronize{
169          instance.cancel(Thread.current, id)
170        }
171      end
172
173      ##
174      # Creates a new TimeoutHandler.  You should use ::register and ::cancel
175      # instead of creating the timeout handler directly.
176      def initialize
177        @timeout_info = Hash.new
178        Thread.start{
179          while true
180            now = Time.now
181            @timeout_info.keys.each{|thread|
182              ary = @timeout_info[thread]
183              next unless ary
184              ary.dup.each{|info|
185                time, exception = *info
186                interrupt(thread, info.object_id, exception) if time < now
187              }
188            }
189            sleep 0.5
190          end
191        }
192      end
193
194      ##
195      # Interrupts the timeout handler +id+ and raises +exception+
196      def interrupt(thread, id, exception)
197        TimeoutMutex.synchronize{
198          if cancel(thread, id) && thread.alive?
199            thread.raise(exception, "execution timeout")
200          end
201        }
202      end
203
204      ##
205      # Registers a new timeout handler
206      #
207      # +time+:: Timeout in seconds
208      # +exception+:: Exception to raise when timeout elapsed
209      def register(thread, time, exception)
210        @timeout_info[thread] ||= Array.new
211        @timeout_info[thread] << [time, exception]
212        return @timeout_info[thread].last.object_id
213      end
214
215      ##
216      # Cancels the timeout handler +id+
217      def cancel(thread, id)
218        if ary = @timeout_info[thread]
219          ary.delete_if{|info| info.object_id == id }
220          if ary.empty?
221            @timeout_info.delete(thread)
222          end
223          return true
224        end
225        return false
226      end
227    end
228
229    ##
230    # Executes the passed block and raises +exception+ if execution takes more
231    # than +seconds+.
232    #
233    # If +seconds+ is zero or nil, simply executes the block
234    def timeout(seconds, exception=Timeout::Error)
235      return yield if seconds.nil? or seconds.zero?
236      # raise ThreadError, "timeout within critical session" if Thread.critical
237      id = TimeoutHandler.register(seconds, exception)
238      begin
239        yield(seconds)
240      ensure
241        TimeoutHandler.cancel(id)
242      end
243    end
244    module_function :timeout
245  end
246end
247