nodejs的纯js版websocket服务器程序

nodejs使得做客户端的人也可以做服务器开发，比较可惜的是目前HTML5客户端支持的WebSocket协议nodejs自身不支持。不过也有很多完善的模块提供使用。

我写这个模块的目的只是为了体验下nodejs的服务器开发而已，打算自己以后写得玩的就用这个算了。其中websocket.js是Websocket服务器模块，server.js是测试代码，跟以前写的一样，模块不支持大于65535的数据包发送和接收。

首先是server.js，用法跟客户端类似：

//引入模块
var server=require("./websocket.js");
//监听"/"请求，可以加入多个监听。
server.addListener("/",function(ws){
	ws.onmessage=function(packet){
		ws.send(new Buffer("Hello World!"));
	}
	ws.onclose=function(){
		console.log("session disconnect");
	}
});
server.keepAlive=3000000;
server.timeout=0;
//开启服务器
server.start("127.0.0.1",8766);

下面是模块的具体实现：

var net=require("net");
var crypto=require('crypto');

String.prototype.trim=function(){return this.replace(/^\s+|\s+$/g, '');};

function WebSocket(socket){
	this._frame={
		step:0,
		fin:false,
		opcode:0,
		mask:false,
		maskingKey:null,
		length: 0,
		current: 0,
		buffer: null
	};
	this._socket=socket;
	this._requestString='';

	this.onopen=null;
	this.onclose=null;
	this.onmessage=null;

	this.handshake=true;
	this.uri="";
	this.method="";
	this.headers=[];
	this.cookies={};
}

WebSocket.prototype={
	send: function(data){
		if(data.length>65535){
			console.error("packet too large.");
			return false;
		}
		var len=data.length+2;
		if(len>=126) len+=2;

		var packet=new Buffer(len);
		len=data.length;
		packet[0]=130;
		var cur=2;
		if(len>=126){
			packet[1]=254;
			packet[2]=(len>>8)&0xff;
			packet[3]=len&0xff;
			cur=4;
		}
		else packet[1]=len;
		for(var i=0;i<len;i++){
			packet[cur+i]=data[i];
		}
		return this._socket.write(packet);
	},
	addBuffer: function(data){
		var frame=this._frame;
		for(var i=0;i<data.length;i++){
			var c=data[i];
			switch(frame.step){
				case 4: //raw data
					frame.buffer[frame.current]=c;
					frame.current++;
					if(frame.current>=frame.length){
						if(frame.mask){
							for(var j=0;j<frame.length;j++){
								frame.buffer[j]=frame.buffer[j] ^ frame.maskingKey[j%4];
							}
						}
						if(this.onmessage) this.onmessage(frame);
						frame.step=0;
					}
					break;
				case 3: //mask
					frame.maskingKey[frame.current]=c;
					frame.current++;
					if(frame.current>=4){
						frame.current=0;
						frame.step++;
					}
					break;
				case 2: //2 bit length
					if(frame.current==0){
						frame.length=c*256;
					}
					else{
						frame.length+=c;
						frame.buffer=new Buffer(frame.length);
						frame.step++;
						frame.current=0;
					}
					break;
				case 1:
					frame.mask=(c&128)==128;
					frame.length=c&127;
					if(frame.mask){
						frame.maskingKey=new Buffer(4);
					}
					if(frame.length<126){
						if(frame.mask) frame.step=3;
						else frame.step=4;
						frame.buffer=new Buffer(frame.length);
					}
					else if(frame.length==126){
						frame.current=0;
						frame.step++;
					}
					else{
						this.send(new Buffer("Server Error: big data ( >65535 bytes ) does not supported."));
						this.close();
						return;
					}
					break;
				case 0:
					frame.current=0;
					frame.fin=(c&128)==128;
					frame.opcode=c&0xf;
					frame.step++;
					break;
			}
		}
	},
	close: function(data){
		if(this.onclose) this.onclose();
		this._socket.end(data);
	},
	addHandshakeBuffer: function(data){
		this._requestString+=data.toString();
		if(this._requestString.length>10240){
			this._requestString=null;
			this.close();
			return -1;
		}
		if(this._requestString.indexOf("\r\n\r\n")!=-1){
			var lines=this._requestString.trim().split("\r\n");
			var t=lines[0].split(/\s+/);
			this.method=t[0];
			this.uri=t[1];
			for(var i=1;i<lines.length;i++){
				var line=lines[i];
				t=line.split(":",2);
				this.headers[t[0].trim()]=t[1].trim();
			}
			var key=this.headers['Sec-WebSocket-Key'];
			if(this.headers["Upgrade"].toLowerCase()!='websocket' || !key){
				this.close();
				return -1;
			}
			var cookie=this.headers['Cookie'];
			if(cookie){
				var cookies=cookie.split(/\s*;\s*/);
				for(var i=0;i<cookies.length;i++){
					t=cookies[i].split(/\s*=\s*/);
					this.cookies[t[0]]=t[1];
				}
			}
			var responseKey=key+"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
			var hash=crypto.createHash('sha1');
			hash.update(responseKey);
			responseKey=hash.digest('base64');
			var response="HTTP/1.1 101 Switching Protocols\r\n"+
				"Upgrade: websocket\r\n"+
				"Connection: Upgrade\r\n" +
				"Sec-WebSocket-Accept: "+responseKey+"\r\n\r\n"
			this._socket.write(response);

			this.handshake=false;
			return 0;
		}
		return 1;
	}
}
module.exports= {
	server: null,
	listeners: {},
	keepAlive: 0,
	timeout: 0,
	addListener: function(uri,callback){
		this.listeners[uri]=callback;
	},
	start: function(ip,port){
		var me=this;
		this.server = net.createServer(function(socket){
			var ws=new WebSocket(socket);
			if(this.keepAlive>0){
				socket.setKeepAlive(true,this.keepAlive);
			}
			if(this.timeout>0){
				socket.setTimeout(this.timeout,function(){
					ws.close();
				});
			}
			socket.addListener('data', function (data) {
				if(ws.handshake){
					var r=ws.addHandshakeBuffer(data);
					if(r==0){
						var callback=me.listeners[ws.uri];
						if(callback) callback(ws);
						else{
							ws.close();
						}
					}
				}
				else{
					ws.addBuffer(data);
				}
			});
			socket.addListener("close",function(hasError){
				console.log("close");
				ws.close();
			});
		});
		this.server.listen(port,ip);
	}
};