diff --git a/tensorboard/components/plugin_util/BUILD b/tensorboard/components/experimental/plugin_util/BUILD similarity index 100% rename from tensorboard/components/plugin_util/BUILD rename to tensorboard/components/experimental/plugin_util/BUILD diff --git a/tensorboard/components/plugin_util/message.ts b/tensorboard/components/experimental/plugin_util/message.ts similarity index 67% rename from tensorboard/components/plugin_util/message.ts rename to tensorboard/components/experimental/plugin_util/message.ts index d9efb9612c..805f2f2752 100644 --- a/tensorboard/components/plugin_util/message.ts +++ b/tensorboard/components/experimental/plugin_util/message.ts @@ -13,6 +13,11 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +/** + * This file defines utilities shared by TensorBoard (plugin host) and the + * dynamic plugin library, used by plugin authors. + */ + export type PayloadType = | null | undefined @@ -27,9 +32,10 @@ export type PayloadType = export interface Message { type: string; - id: string; + id: number; payload: PayloadType; error: string | null; + isReply: boolean; } export type MessageType = string; @@ -40,21 +46,13 @@ interface PromiseResolver { reject: (error: Error) => void; } -export abstract class IPC { - private idPrefix: string; +export class IPC { private id = 0; - private readonly responseWaits = new Map(); + 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(''); + constructor(private port: MessagePort) { + this.port.addEventListener('message', (event) => this.onMessage(event)); } listen(type: MessageType, callback: MessageCallback) { @@ -66,13 +64,11 @@ export abstract class IPC { } 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)) { + if (message.isReply) { + if (!this.responseWaits.has(message.id)) return; const {id, payload, error} = message; const {resolve, reject} = this.responseWaits.get(id); this.responseWaits.delete(id); @@ -100,22 +96,19 @@ export abstract class IPC { id: message.id, payload, error, + isReply: true, }; - this.postMessage(event.source, JSON.stringify(replyMessage)); + this.postMessage(replyMessage); } - private postMessage(targetWindow: Window, message: string) { - targetWindow.postMessage(message, '*'); + private postMessage(message: Message) { + this.port.postMessage(JSON.stringify(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)); + sendMessage(type: MessageType, payload: PayloadType): Promise { + const id = this.id++; + const message: Message = {type, id, payload, error: null, isReply: false}; + this.postMessage(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/experimental/plugin_util/plugin-guest.html similarity index 100% rename from tensorboard/components/plugin_util/plugin-guest.html rename to tensorboard/components/experimental/plugin_util/plugin-guest.html diff --git a/tensorboard/components/plugin_util/plugin-guest.ts b/tensorboard/components/experimental/plugin_util/plugin-guest.ts similarity index 70% rename from tensorboard/components/plugin_util/plugin-guest.ts rename to tensorboard/components/experimental/plugin_util/plugin-guest.ts index 9557c92a10..5837473518 100644 --- a/tensorboard/components/plugin_util/plugin-guest.ts +++ b/tensorboard/components/experimental/plugin_util/plugin-guest.ts @@ -12,19 +12,25 @@ 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); - } +import {IPC, Message} from './message.js'; + +/** + * This code is part of a public bundle provided to plugin authors, + * and runs within an IFrame to setup communication with TensorBoard's frame. + */ +if (!window.parent) { + throw Error('The library must run within a TensorBoard iframe-based plugin.'); } +const channel = new MessageChannel(); +const ipc = new IPC(channel.port1); +channel.port1.start(); + +const VERSION = 'experimental'; +window.parent.postMessage(`${VERSION}.bootstrap`, '*', [channel.port2]); + // Only export for testability. -export const _guestIPC = new GuestIPC(); +export const _guestIPC = ipc; /** * Sends a message to the parent frame. diff --git a/tensorboard/components/plugin_util/plugin-host.html b/tensorboard/components/experimental/plugin_util/plugin-host.html similarity index 100% rename from tensorboard/components/plugin_util/plugin-host.html rename to tensorboard/components/experimental/plugin_util/plugin-host.html diff --git a/tensorboard/components/experimental/plugin_util/plugin-host.ts b/tensorboard/components/experimental/plugin_util/plugin-host.ts new file mode 100644 index 0000000000..1de5c554cb --- /dev/null +++ b/tensorboard/components/experimental/plugin_util/plugin-host.ts @@ -0,0 +1,100 @@ +/* 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, MessageType, PayloadType, MessageCallback} from './message.js'; + +const portIPCs = new Set(); +const VERSION = 'experimental'; +const listeners = new Map(); + +// TODO(@psybuzz): replace this and the port cleanup logic in broadcast() with +// a MutationObserver to notify us when iframes disconnect. +const ipcToFrame = new Map(); + +// The initial Window-level listener is needed to bootstrap only. +// All further communication is done over MessagePorts. +window.addEventListener('message', (event) => { + if (event.data !== `${VERSION}.bootstrap`) return; + const port = event.ports[0]; + if (!port) return; + const frame = event.source ? event.source.frameElement : null; + if (!frame) return; + onBootstrap(port, frame as HTMLIFrameElement); +}); + +function onBootstrap(port: MessagePort, frame: HTMLIFrameElement) { + const portIPC = new IPC(port); + portIPCs.add(portIPC); + ipcToFrame.set(portIPC, frame); + port.start(); + + for (const [type, callback] of listeners) { + portIPC.listen(type, callback); + } +} + +function _broadcast( + type: MessageType, + payload: PayloadType +): Promise { + for (const ipc of portIPCs) { + if (!ipcToFrame.get(ipc).isConnected) { + portIPCs.delete(ipc); + ipcToFrame.delete(ipc); + } + } + + const promises = [...portIPCs].map((ipc) => ipc.sendMessage(type, payload)); + return Promise.all(promises); +} + +function _listen(type: MessageType, callback: MessageCallback) { + listeners.set(type, callback); + for (const ipc of portIPCs) { + ipc.listen(type, callback); + } +} + +function _unlisten(type: MessageType) { + listeners.delete(type); + for (const ipc of portIPCs) { + ipc.unlisten(type); + } +} + +export const broadcast = _broadcast; +export const listen = _listen; +export const unlisten = _unlisten; + +namespace tf_plugin { + /** + * Sends a message to all frames. Individual frames decide whether or not to + * listen. + * @return Promise that resolves with a list of payloads from each plugin's + * response (or null) to the message. + * + * @example + * const someList = await broadcast('v1.some.type.guest.understands'); + * // do fun things with someList. + */ + export const broadcast = _broadcast; + /** + * Subscribes to messages of a type specified for all frames. + */ + export const listen = _listen; + /** + * Unsubscribes to messages of a type specified for all frames. + */ + export const unlisten = _unlisten; +} // namespace tf_plugin diff --git a/tensorboard/components/plugin_util/test/BUILD b/tensorboard/components/experimental/plugin_util/test/BUILD similarity index 88% rename from tensorboard/components/plugin_util/test/BUILD rename to tensorboard/components/experimental/plugin_util/test/BUILD index 0de10f7665..eea69f1d29 100644 --- a/tensorboard/components/plugin_util/test/BUILD +++ b/tensorboard/components/experimental/plugin_util/test/BUILD @@ -10,8 +10,8 @@ licenses(["notice"]) # Apache 2.0 tf_web_test( name = "test", + src = "/tf-plugin/test/test_binary.html", web_library = ":test_web_library", - src = "/tf-plugin/test/test_binary.html" ) # HACK: specifying tensorboard_html_binary on tf_web_test causes certain @@ -53,8 +53,8 @@ tf_web_library( ], path = "/tf-plugin/test", deps = [ - "//tensorboard/components/plugin_util:plugin_host", - "//tensorboard/components/plugin_util:plugin_guest", + "//tensorboard/components/experimental/plugin_util:plugin_guest", + "//tensorboard/components/experimental/plugin_util:plugin_host", "//tensorboard/components/tf_imports:web_component_tester", ], ) diff --git a/tensorboard/components/plugin_util/test/iframe.html b/tensorboard/components/experimental/plugin_util/test/iframe.html similarity index 100% rename from tensorboard/components/plugin_util/test/iframe.html rename to tensorboard/components/experimental/plugin_util/test/iframe.html diff --git a/tensorboard/components/plugin_util/test/iframe.ts b/tensorboard/components/experimental/plugin_util/test/iframe.ts similarity index 100% rename from tensorboard/components/plugin_util/test/iframe.ts rename to tensorboard/components/experimental/plugin_util/test/iframe.ts diff --git a/tensorboard/components/plugin_util/test/plugin-test.ts b/tensorboard/components/experimental/plugin_util/test/plugin-test.ts similarity index 75% rename from tensorboard/components/plugin_util/test/plugin-test.ts rename to tensorboard/components/experimental/plugin_util/test/plugin-test.ts index 4c1500ee2f..b56a30f1d8 100644 --- a/tensorboard/components/plugin_util/test/plugin-test.ts +++ b/tensorboard/components/experimental/plugin_util/test/plugin-test.ts @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ import * as pluginHost from '../plugin-host.js'; +import {Message} from '../message.js'; namespace tf_plugin.test { const {expect} = chai; @@ -54,25 +55,29 @@ namespace tf_plugin.test { { spec: 'host (src) to guest (dest)', beforeEachFunc: function() { + this.destWindow = this.guestWindow; this.destListen = this.guestWindow.test.listen; this.destUnlisten = this.guestWindow.test.unlisten; + this.destSendMessage = this.guestWindow.test.sendMessage; this.srcSendMessage = (type, payload) => { - return pluginHost.sendMessage(this.guestFrame, type, payload); + return pluginHost + .broadcast(type, payload) + .then(([result]) => result); }; - this.destPostMessageSpy = () => - this.sandbox.spy(this.guestWindow.test._guestIPC, 'postMessage'); }, }, { spec: 'guest (src) to host (dest)', beforeEachFunc: function() { + this.destWindow = window; this.destListen = pluginHost.listen; this.destUnlisten = pluginHost.unlisten; - this.srcSendMessage = (type, payload) => { - return this.guestWindow.test.sendMessage(type, payload); + this.destSendMessage = (type, payload) => { + return pluginHost + .broadcast(type, payload) + .then(([result]) => result); }; - this.destPostMessageSpy = () => - this.sandbox.spy(pluginHost._hostIPC, 'postMessage'); + this.srcSendMessage = this.guestWindow.test.sendMessage; }, }, ].forEach(({spec, beforeEachFunc}) => { @@ -105,15 +110,12 @@ namespace tf_plugin.test { }); 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']); }); @@ -175,6 +177,46 @@ namespace tf_plugin.test { expect(barCb.callCount).to.equal(1); }); + + it('ignores foreign postMessages', async function() { + const barCb = this.sandbox.stub(); + this.destListen('bar', barCb); + const fakeMessage: Message = { + type: 'bar', + id: 0, + payload: '', + error: null, + isReply: false, + }; + this.destWindow.postMessage(JSON.stringify(fakeMessage), '*'); + + // Await another message to ensure fake message was handled in dest. + await this.srcSendMessage('not-bar'); + expect(barCb).to.not.have.been.called; + }); + + it('processes messages while waiting for a reponse', async function() { + let resolveLongTask = null; + this.destListen('longTask', () => { + return new Promise((resolve) => { + resolveLongTask = resolve; + }); + }); + + const longTaskStub = this.sandbox.stub(); + const longTaskPromise = this.srcSendMessage('longTask', 'hello').then( + longTaskStub + ); + + await this.srcSendMessage('foo'); + await this.destSendMessage('bar'); + expect(longTaskStub).to.not.have.been.called; + + resolveLongTask('payload'); + const longTaskResult = await longTaskPromise; + expect(longTaskStub).to.have.been.calledOnce; + expect(longTaskStub).to.have.been.calledWith('payload'); + }); }); }); }); diff --git a/tensorboard/components/plugin_util/test/tests.html b/tensorboard/components/experimental/plugin_util/test/tests.html similarity index 100% rename from tensorboard/components/plugin_util/test/tests.html rename to tensorboard/components/experimental/plugin_util/test/tests.html diff --git a/tensorboard/components/plugin_util/plugin-host.ts b/tensorboard/components/plugin_util/plugin-host.ts deleted file mode 100644 index 48da663013..0000000000 --- a/tensorboard/components/plugin_util/plugin-host.ts +++ /dev/null @@ -1,57 +0,0 @@ -/* 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