diff --git a/tensorboard/components/plugin_util/BUILD b/tensorboard/components/plugin_util/BUILD new file mode 100644 index 0000000000..622dfeef6e --- /dev/null +++ b/tensorboard/components/plugin_util/BUILD @@ -0,0 +1,39 @@ +package(default_visibility = ["//tensorboard:internal"]) + +load("//tensorboard/defs:web.bzl", "tf_web_library") + +licenses(["notice"]) # Apache 2.0 + +tf_web_library( + name = "plugin_host", + srcs = [ + "plugin-host.html", + "plugin-host.ts", + ], + path = "/tf-plugin", + deps = [ + ":plugin_lib", + "//tensorboard/components/tf_backend", + ], +) + +tf_web_library( + name = "plugin_lib", + srcs = [ + "message.ts", + ], + path = "/tf-plugin", +) + +tf_web_library( + name = "plugin_guest", + srcs = [ + "plugin-guest.html", + "plugin-guest.ts", + ], + path = "/tf-plugin", + visibility = ["//visibility:public"], + deps = [ + ":plugin_lib", + ], +) diff --git a/tensorboard/components/plugin_util/message.ts b/tensorboard/components/plugin_util/message.ts new file mode 100644 index 0000000000..d9efb9612c --- /dev/null +++ b/tensorboard/components/plugin_util/message.ts @@ -0,0 +1,123 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +export type PayloadType = + | null + | undefined + | string + | string[] + | boolean + | boolean[] + | number + | number[] + | object + | object[]; + +export interface Message { + type: string; + id: string; + payload: PayloadType; + error: string | null; +} + +export type MessageType = string; +export type MessageCallback = (payload: any) => any; + +interface PromiseResolver { + resolve: (data: any) => void; + reject: (error: Error) => void; +} + +export abstract class IPC { + private idPrefix: string; + private id = 0; + private readonly responseWaits = new Map(); + private readonly listeners = new Map(); + + constructor() { + window.addEventListener('message', this.onMessage.bind(this)); + + // TODO(tensorboard-team): remove this by using MessageChannel. + const randomArray = new Uint8Array(16); + window.crypto.getRandomValues(randomArray); + this.idPrefix = Array.from(randomArray) + .map((int: number) => int.toString(16)) + .join(''); + } + + listen(type: MessageType, callback: MessageCallback) { + this.listeners.set(type, callback); + } + + unlisten(type: MessageType) { + this.listeners.delete(type); + } + + private async onMessage(event: MessageEvent) { + // There are instances where random browser extensions send messages. + if (typeof event.data !== 'string') return; + + const message = JSON.parse(event.data) as Message; + const callback = this.listeners.get(message.type); + + if (this.responseWaits.has(message.id)) { + const {id, payload, error} = message; + const {resolve, reject} = this.responseWaits.get(id); + this.responseWaits.delete(id); + if (error) { + reject(new Error(error)); + } else { + resolve(payload); + } + return; + } + + let payload = null; + let error = null; + if (this.listeners.has(message.type)) { + const callback = this.listeners.get(message.type); + try { + const result = await callback(message.payload); + payload = result; + } catch (e) { + error = e; + } + } + const replyMessage: Message = { + type: message.type, + id: message.id, + payload, + error, + }; + this.postMessage(event.source, JSON.stringify(replyMessage)); + } + + private postMessage(targetWindow: Window, message: string) { + targetWindow.postMessage(message, '*'); + } + + protected sendMessageToWindow( + targetWindow: Window, + type: MessageType, + payload: PayloadType + ): Promise { + const id = `${this.idPrefix}_${this.id++}`; + const message: Message = {type, id, payload, error: null}; + this.postMessage(targetWindow, JSON.stringify(message)); + return new Promise((resolve, reject) => { + this.responseWaits.set(id, {resolve, reject}); + }); + } +} diff --git a/tensorboard/components/plugin_util/plugin-guest.html b/tensorboard/components/plugin_util/plugin-guest.html new file mode 100644 index 0000000000..807d96f6e7 --- /dev/null +++ b/tensorboard/components/plugin_util/plugin-guest.html @@ -0,0 +1,18 @@ + + + diff --git a/tensorboard/components/plugin_util/plugin-guest.ts b/tensorboard/components/plugin_util/plugin-guest.ts new file mode 100644 index 0000000000..9557c92a10 --- /dev/null +++ b/tensorboard/components/plugin_util/plugin-guest.ts @@ -0,0 +1,47 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +import {IPC, Message, MessageType, PayloadType} from './message.js'; + +class GuestIPC extends IPC { + /** + * payload must be JSON serializable. + */ + sendMessage(type: MessageType, payload: PayloadType): Promise { + return this.sendMessageToWindow(window.parent, type, payload); + } +} + +// Only export for testability. +export const _guestIPC = new GuestIPC(); + +/** + * Sends a message to the parent frame. + * @return Promise that resolves with a payload from parent in response to this message. + * + * @example + * const someList = await sendMessage('v1.some.type.parent.understands'); + * // do fun things with someList. + */ +export const sendMessage = _guestIPC.sendMessage.bind(_guestIPC); + +/** + * Subscribes a callback to a message with particular type. + */ +export const listen = _guestIPC.listen.bind(_guestIPC); + +/** + * Unsubscribes a callback to a message. + */ +export const unlisten = _guestIPC.unlisten.bind(_guestIPC); diff --git a/tensorboard/components/plugin_util/plugin-host.html b/tensorboard/components/plugin_util/plugin-host.html new file mode 100644 index 0000000000..d1b6f3ae89 --- /dev/null +++ b/tensorboard/components/plugin_util/plugin-host.html @@ -0,0 +1,20 @@ + + + + + diff --git a/tensorboard/components/plugin_util/plugin-host.ts b/tensorboard/components/plugin_util/plugin-host.ts new file mode 100644 index 0000000000..48da663013 --- /dev/null +++ b/tensorboard/components/plugin_util/plugin-host.ts @@ -0,0 +1,57 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +import {IPC, Message, MessageType, PayloadType} from './message.js'; + +class HostIPC extends IPC { + sendMessage( + iframe: HTMLIFrameElement, + type: MessageType, + payload: PayloadType + ): Promise { + return this.sendMessageToWindow(iframe.contentWindow, type, payload); + } +} + +const hostIPC = new HostIPC(); +const _listen = hostIPC.listen.bind(hostIPC); +const _unlisten = hostIPC.unlisten.bind(hostIPC); +const _sendMessage = hostIPC.sendMessage.bind(hostIPC); + +export const sendMessage = _sendMessage; +export const listen = _listen; +export const unlisten = _unlisten; + +// Export for testability. +export const _hostIPC = hostIPC; + +namespace tf_plugin { + /** + * Sends a message to the frame specified. + * @return Promise that resolves with a payload from frame in response to the message. + * + * @example + * const someList = await sendMessage('v1.some.type.guest.understands'); + * // do fun things with someList. + */ + export const sendMessage = _sendMessage; + /** + * Subscribes to messages from specified frame of a type specified. + */ + export const listen = _listen; + /** + * Unsubscribes to messages from specified frame of a type specified. + */ + export const unlisten = _unlisten; +} // namespace tf_plugin diff --git a/tensorboard/components/plugin_util/test/BUILD b/tensorboard/components/plugin_util/test/BUILD new file mode 100644 index 0000000000..0de10f7665 --- /dev/null +++ b/tensorboard/components/plugin_util/test/BUILD @@ -0,0 +1,60 @@ +package( + default_testonly = True, + default_visibility = ["//tensorboard:internal"], +) + +load("//tensorboard/defs:web.bzl", "tf_web_library", "tf_web_test") +load("//tensorboard/defs:vulcanize.bzl", "tensorboard_html_binary") + +licenses(["notice"]) # Apache 2.0 + +tf_web_test( + name = "test", + web_library = ":test_web_library", + src = "/tf-plugin/test/test_binary.html" +) + +# HACK: specifying tensorboard_html_binary on tf_web_test causes certain +# environment to throw exception but wrapping tensorbard_html_binary with +# tf_web_library seems to be okay. +tf_web_library( + name = "test_web_library", + srcs = [ + ":test_binary.html", + ], + path = "/tf-plugin/test", + deps = [ + ":test_lib", + "//tensorboard/components/tf_imports:web_component_tester", + ], +) + +tensorboard_html_binary( + name = "test_binary", + # Disable advanced optimization to prevent check for WebComponentTester + # gobals. + compilation_level = "SIMPLE", + # Requires for compiling `import`s away. + compile = True, + input_path = "/tf-plugin/test/tests.html", + output_path = "/tf-plugin/test/test_binary.html", + deps = [ + ":test_lib", + ], +) + +tf_web_library( + name = "test_lib", + srcs = [ + "iframe.html", + "iframe.ts", + "plugin-test.ts", + "tests.html", + ], + path = "/tf-plugin/test", + deps = [ + "//tensorboard/components/plugin_util:plugin_host", + "//tensorboard/components/plugin_util:plugin_guest", + "//tensorboard/components/tf_imports:web_component_tester", + ], +) diff --git a/tensorboard/components/plugin_util/test/iframe.html b/tensorboard/components/plugin_util/test/iframe.html new file mode 100644 index 0000000000..9aff72251d --- /dev/null +++ b/tensorboard/components/plugin_util/test/iframe.html @@ -0,0 +1,20 @@ + + + + + diff --git a/tensorboard/components/plugin_util/test/iframe.ts b/tensorboard/components/plugin_util/test/iframe.ts new file mode 100644 index 0000000000..c1c979687c --- /dev/null +++ b/tensorboard/components/plugin_util/test/iframe.ts @@ -0,0 +1,19 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +import {sendMessage, listen, unlisten, _guestIPC} from '../plugin-guest.js'; + +const win = window as any; +win.test = {sendMessage, listen, unlisten, _guestIPC}; diff --git a/tensorboard/components/plugin_util/test/plugin-test.ts b/tensorboard/components/plugin_util/test/plugin-test.ts new file mode 100644 index 0000000000..4c1500ee2f --- /dev/null +++ b/tensorboard/components/plugin_util/test/plugin-test.ts @@ -0,0 +1,181 @@ +/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +import * as pluginHost from '../plugin-host.js'; + +namespace tf_plugin.test { + const {expect} = chai; + const template = document.getElementById( + 'iframe-template' + ) as HTMLTemplateElement; + + describe('plugin-util', () => { + beforeEach(function(done) { + const iframeFrag = document.importNode(template.content, true); + const iframe = iframeFrag.firstElementChild as HTMLIFrameElement; + document.body.appendChild(iframe); + this.guestFrame = iframe; + this.guestWindow = iframe.contentWindow; + // Must wait for the JavaScript to be loaded on the child frame. + this.guestWindow.addEventListener('load', () => done()); + + this.sandbox = sinon.sandbox.create(); + }); + + afterEach(function() { + document.body.removeChild(this.guestFrame); + this.sandbox.restore(); + }); + + it('setUp sanity check', function() { + expect(this.guestWindow.test) + .to.have.property('sendMessage') + .that.is.a('function'); + expect(this.guestWindow.test) + .to.have.property('listen') + .that.is.a('function'); + expect(this.guestWindow.test) + .to.have.property('unlisten') + .that.is.a('function'); + }); + + [ + { + spec: 'host (src) to guest (dest)', + beforeEachFunc: function() { + this.destListen = this.guestWindow.test.listen; + this.destUnlisten = this.guestWindow.test.unlisten; + this.srcSendMessage = (type, payload) => { + return pluginHost.sendMessage(this.guestFrame, type, payload); + }; + this.destPostMessageSpy = () => + this.sandbox.spy(this.guestWindow.test._guestIPC, 'postMessage'); + }, + }, + { + spec: 'guest (src) to host (dest)', + beforeEachFunc: function() { + this.destListen = pluginHost.listen; + this.destUnlisten = pluginHost.unlisten; + this.srcSendMessage = (type, payload) => { + return this.guestWindow.test.sendMessage(type, payload); + }; + this.destPostMessageSpy = () => + this.sandbox.spy(pluginHost._hostIPC, 'postMessage'); + }, + }, + ].forEach(({spec, beforeEachFunc}) => { + describe(spec, () => { + beforeEach(beforeEachFunc); + + beforeEach(function() { + this.onMessage = this.sandbox.stub(); + this.destListen('messageType', this.onMessage); + }); + + it('sends a message to dest', async function() { + await this.srcSendMessage('messageType', 'hello'); + expect(this.onMessage.callCount).to.equal(1); + expect(this.onMessage.firstCall.args).to.deep.equal(['hello']); + }); + + it('sends a message a random payload not by ref', async function() { + const payload = { + foo: 'foo', + bar: { + baz: 'baz', + }, + }; + await this.srcSendMessage('messageType', payload); + expect(this.onMessage.callCount).to.equal(1); + + expect(this.onMessage.firstCall.args[0]).to.not.equal(payload); + expect(this.onMessage.firstCall.args[0]).to.deep.equal(payload); + }); + + it('resolves when dest replies with ack', async function() { + const destPostMessage = this.destPostMessageSpy(); + const sendMessageP = this.srcSendMessage('messageType', 'hello'); + + expect(this.onMessage.callCount).to.equal(0); + expect(destPostMessage.callCount).to.equal(0); + + await sendMessageP; + expect(this.onMessage.callCount).to.equal(1); + expect(destPostMessage.callCount).to.equal(1); + expect(this.onMessage.firstCall.args).to.deep.equal(['hello']); + }); + + it('triggers, on dest, a cb for the matching type', async function() { + const barCb = this.sandbox.stub(); + this.destListen('bar', barCb); + + await this.srcSendMessage('bar', 'soap'); + + expect(this.onMessage.callCount).to.equal(0); + expect(barCb.callCount).to.equal(1); + expect(barCb.firstCall.args).to.deep.equal(['soap']); + }); + + it('supports single listener for a type', async function() { + const barCb1 = this.sandbox.stub(); + const barCb2 = this.sandbox.stub(); + this.destListen('bar', barCb1); + this.destListen('bar', barCb2); + + await this.srcSendMessage('bar', 'soap'); + + expect(barCb1.callCount).to.equal(0); + expect(barCb2.callCount).to.equal(1); + expect(barCb2.firstCall.args).to.deep.equal(['soap']); + }); + + describe('dest message handling', () => { + [ + {specName: 'undefined', payload: null, expectDeep: false}, + {specName: 'null', payload: undefined, expectDeep: false}, + {specName: 'string', payload: 'something', expectDeep: false}, + {specName: 'number', payload: 3.14, expectDeep: false}, + {specName: 'object', payload: {some: 'object'}, expectDeep: true}, + {specName: 'array', payload: ['a', 'b', 'c'], expectDeep: true}, + ].forEach(({specName, payload, expectDeep}) => { + it(specName, async function() { + this.destListen('bar', () => payload); + + const response = await this.srcSendMessage('bar', 'soap'); + + if (expectDeep) { + expect(response).to.deep.equal(payload); + } else { + expect(response).to.equal(payload); + } + }); + }); + }); + + it('unregister a callback with unlisten', async function() { + const barCb = this.sandbox.stub(); + this.destListen('bar', barCb); + await this.srcSendMessage('bar', 'soap'); + expect(barCb.callCount).to.equal(1); + this.destUnlisten('bar'); + + await this.srcSendMessage('bar', 'soap'); + + expect(barCb.callCount).to.equal(1); + }); + }); + }); + }); +} // namespace tf_plugin.test diff --git a/tensorboard/components/plugin_util/test/tests.html b/tensorboard/components/plugin_util/test/tests.html new file mode 100644 index 0000000000..76946a9b0e --- /dev/null +++ b/tensorboard/components/plugin_util/test/tests.html @@ -0,0 +1,23 @@ + + + + + +