websocket-driver-ruby/ext/websocket-driver/WebsocketMaskService.java

package com.jcoglan.websocket;

import java.io.IOException;
import org.jruby.Ruby;
import org.jruby.RubyClass;
import org.jruby.RubyModule;
import org.jruby.RubyObject;
import org.jruby.RubyString;
import org.jruby.anno.JRubyMethod;
import org.jruby.runtime.ObjectAllocator;
import org.jruby.runtime.ThreadContext;
import org.jruby.runtime.builtin.IRubyObject;
import org.jruby.runtime.load.BasicLibraryService;

public class WebsocketMaskService implements BasicLibraryService {
    private Ruby runtime;

    public boolean basicLoad(Ruby runtime) throws IOException {
        this.runtime = runtime;

        RubyModule websocket = runtime.defineModule("WebSocket");
        RubyClass webSocketMask = websocket.defineClassUnder("Mask", runtime.getObject(), getAllocator());

        webSocketMask.defineAnnotatedMethods(WebsocketMask.class);
        return true;
    }

    ObjectAllocator getAllocator() {
        return new ObjectAllocator() {
            public IRubyObject allocate(Ruby runtime, RubyClass rubyClass) {
                return new WebsocketMask(runtime, rubyClass);
            }
        };
    }

    public class WebsocketMask extends RubyObject {
        public WebsocketMask(final Ruby runtime, RubyClass rubyClass) {
            super(runtime, rubyClass);
        }

        @JRubyMethod
        public IRubyObject mask(ThreadContext context, IRubyObject payload, IRubyObject mask) {
            if (mask.isNil()) return payload;

            byte[] payload_a = ((RubyString)payload).getBytes();
            byte[] mask_a    = ((RubyString)mask).getBytes();
            int i, n         = payload_a.length;

            if (n == 0) return payload;

            for (i = 0; i < n; i++) {
                payload_a[i] ^= mask_a[i % 4];
            }
            return RubyString.newStringNoCopy(runtime, payload_a);
        }
    }
}