lib/nyoibo/daemon.rb in nyoibo-0.0.2 vs lib/nyoibo/daemon.rb in nyoibo-0.0.3

@@ -10,11 +10,11 @@
         EventMachine::WebSocket.start(:host => config.host, :port => config.port) do |ws|
           ws.onopen{
             ws.send "OK Ready"
           }
           ws.onclose{
-            ws.send "OK Bye"
+            ws.close_websocket
           }
           ws.onmessage{|msg|
             @binary ||= ""
             @binary_type ||= 'binary'
             case msg
@@ -22,19 +22,24 @@
               if @binary_type == 'base64'
                 @binary = Base64.decode64(@binary)
               elsif @binary.encoding == Encoding::UTF_8
                 @binary = @binary.unpack('U*').pack('c*')
               end
-              Nyoibo.run_callback(ws.request["path"], @json, @binary)
-              ws.close_connection
+              Nyoibo.run_callback(:after, ws.request["path"], @json, @binary)
+              @binary = @json = nil
+              ws.send("OK Bye")
             when CMD_JSON
               msg.gsub!(CMD_JSON, '')
               @json = JSON.parse(msg)
               @json['size'] = @json['size'].to_i
-              ws.send("NEXT")
+              if Nyoibo.run_callback(:before, ws.request["path"], @json, @binary) == false
+                ws.send("ABORT")
+                ws.close_websocket
+              else
+                ws.send("NEXT")
+              end
             else
-
               if msg =~ TYPE_BASE64
                 msg.gsub!(TYPE_BASE64, '')
                 @binary_type = 'base64'
               end
 
@@ -50,9 +55,12 @@
       }
       if ENV["NYOIBO_ENV"] == "test"
         daemon.call
       else
         @pid = Process.fork &daemon
+        at_exit do
+          Process.kill(:KILL, Nyoibo.pid) if Nyoibo.pid
+        end
       end
     end
   end
 end