# Microsft WinSock2 basic binding module using DL
# by Yoann Guillot, 03-2006

require 'dl'
require 'dl/struct'
require 'dl/import'

module WS2_32
	extend DL::Importable
	dlload 'ws2_32'

	AF_INET = 2
	SOCK_STREAM = 1
	IPPROTO_TCP = 6
	INADDR_ANY = 0
	NO_ERROR = 0
	INVALID_SOCKET = ~0

	typealias 'word', 'unsigned short'
	typealias 'SOCKET', 'long'	# 'unsigned int *'
	
	WSAData = struct ['word version', 'word hversion', 'char desc[257]', 'char sysstat[129]']
	SockAddr_in = struct ['short sin_family', 'short sin_port', 'long sin_addr', 'long zero0', 'long zero1']
	
	extern 'int WSAStartup(unsigned short, wsadata *)'
	extern 'int WSACleanup()'
	extern 'int WSAGetLastError()'
	
	extern 'SOCKET socket(int, int, int)'
	extern 'int closesocket(SOCKET)'
	extern 'int connect(SOCKET, struct sockaddr *, int)'
	extern 'int bind(SOCKET, struct sockaddr *, int)'
	extern 'int listen(SOCKET, int)'
	extern 'SOCKET accept(SOCKET, struct sockaddr *, int *)'
	extern 'int recv(SOCKET, void *, int, int)'
	extern 'int send(SOCKET, void *, int, int)'

	extern 'unsigned short ntohs(unsigned short)'
	extern 'unsigned long ntohl(unsigned long)'
	extern 'char * inet_ntoa(unsigned long)'
	extern 'unsigned long inet_addr(const char *)'

	def self.init
		data = WSAData.malloc
		ret = wSAStartup(0x0202, data)
		raise "Error #{ret} at WSAStartup" if ret != NO_ERROR
		puts "#{data.inspect}" if $VERBOSE
		
		if block_given?
			begin
				yield
			ensure
				self.exit
			end
		else
			true
		end
	end
	
	def self.exit
		wSACleanup
	end

	class Socket
		class << self
			def new(*args)
				o = super
				if block_given?
					begin
						yield o
					ensure
						o.close
						nil
					end
				else
					o
				end
			end
		end
		
		attr_reader :s, :peer
		
		def error(function, code = wSAGetLastError)
			raise "Error #{code} in #{function}"
		end
		
		def initialize(af=AF_INET, type=SOCK_STREAM, proto=IPPROTO_TCP, s=nil, peer=nil)
			@family = af
			@type = type
			@proto = proto
			if s
				@s = s
				@peer = peer
			else
				@s = WS2_32.socket af, type, proto
				error 'socket' if @s == INVALID_SOCKET
			end
		end

		def connect(ip, port, family=@family)
			sa = SockAddr_in.malloc
			sa.sin_family = @family
			sa.sin_addr = WS2_32.inet_addr ip
			sa.sin_port = WS2_32.ntohs port
			r = WS2_32.connect @s, sa, sa.size
			error 'connect' if r != 0
		end
		
		def bind(ip=nil, port=0, family=@family)
			sa = SockAddr_in.malloc
			sa.sin_family = family
			sa.sin_addr = ip ? WS2_32.inet_addr(ip) : INADDR_ANY
			sa.sin_port = WS2_32.ntohs port
			r = WS2_32.bind @s, sa, sa.size
			error 'bind' if r != 0
		end

		def listen(backlog=4)
			r = WS2_32.listen @s, backlog
			error 'listen' if r != 0
		end

		def accept &block
			ia = SockAddr_in.malloc
			il = DL.malloc(DL.sizeof('i'))
			il.struct!('I', 'len')
			il['len'] = ia.size
			a = WS2_32.accept @s, ia, il
			error 'accept' if a == INVALID_SOCKET
			self.class.new @family, @type, @proto, a, ia, &block
		end

		def recv(len = 512, flags = 0)
			buf = DL.malloc len
			r = WS2_32.recv @s, buf, len, flags
			error 'recv' if r < 0
			# nil if r == 0
			buf.to_s[0..r] if r > 0
		end

		def send(str, flags = 0)
			r = WS2_32.send @s, str, str.length, flags
			error 'send' if r <= 0
			r
		end

		def close
			WS2_32.closesocket(@s)
		end
	end
end

if __FILE__ == $0
WS2_32.init {
	WS2_32::Socket.new { |s|
		s.bind("127.0.0.1", 27015)
		s.listen
		s.accept { |a|
			puts "connection from #{WS2_32.inet_ntoa a.peer.sin_addr}:#{WS2_32.ntohs a.peer.sin_port}"
			
			str = a.recv
			puts "received #{str.inspect}"
			r = a.send(str)
			puts "resent #{r} bytes"
		}
	}
}

puts 'hit enter to finish' ; gets
end