diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index 071417b..6697b60 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -517,6 +517,16 @@
         return -1;
     }
 
+    private void setVirtualDisplaySurfaceInternal(IBinder appToken, Surface surface) {
+        synchronized (mSyncRoot) {
+            if (mVirtualDisplayAdapter == null) {
+                return;
+            }
+
+            mVirtualDisplayAdapter.setVirtualDisplaySurfaceLocked(appToken, surface);
+        }
+    }
+
     private void releaseVirtualDisplayInternal(IBinder appToken) {
         synchronized (mSyncRoot) {
             if (mVirtualDisplayAdapter == null) {
@@ -1221,9 +1231,6 @@
                 throw new IllegalArgumentException("width, height, and densityDpi must be "
                         + "greater than 0");
             }
-            if (surface == null) {
-                throw new IllegalArgumentException("surface must not be null");
-            }
             if (callingUid != Process.SYSTEM_UID &&
                     (flags & DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC) != 0) {
                 if (mContext.checkCallingPermission(android.Manifest.permission.CAPTURE_VIDEO_OUTPUT)
@@ -1255,6 +1262,16 @@
         }
 
         @Override // Binder call
+        public void setVirtualDisplaySurface(IBinder appToken, Surface surface) {
+            final long token = Binder.clearCallingIdentity();
+            try {
+                setVirtualDisplaySurfaceInternal(appToken, surface);
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
+        }
+
+        @Override // Binder call
         public void releaseVirtualDisplay(IBinder appToken) {
             final long token = Binder.clearCallingIdentity();
             try {
diff --git a/services/core/java/com/android/server/display/VirtualDisplayAdapter.java b/services/core/java/com/android/server/display/VirtualDisplayAdapter.java
index 95ca0d2..a165f26 100644
--- a/services/core/java/com/android/server/display/VirtualDisplayAdapter.java
+++ b/services/core/java/com/android/server/display/VirtualDisplayAdapter.java
@@ -69,6 +69,13 @@
         return device;
     }
 
+    public void setVirtualDisplaySurfaceLocked(IBinder appToken, Surface surface) {
+        VirtualDisplayDevice device = mVirtualDisplayDevices.get(appToken);
+        if (device != null) {
+            device.setSurfaceLocked(surface);
+        }
+    }
+
     public DisplayDevice releaseVirtualDisplayLocked(IBinder appToken) {
         VirtualDisplayDevice device = mVirtualDisplayDevices.remove(appToken);
         if (device != null) {
@@ -144,6 +151,17 @@
             }
         }
 
+        public void setSurfaceLocked(Surface surface) {
+            if (mSurface != surface) {
+                if ((mSurface != null) != (surface != null)) {
+                    sendDisplayDeviceEventLocked(this, DISPLAY_DEVICE_EVENT_CHANGED);
+                }
+                sendTraversalRequestLocked();
+                mSurface = surface;
+                mInfo = null;
+            }
+        }
+
         @Override
         public DisplayDeviceInfo getDisplayDeviceInfoLocked() {
             if (mInfo == null) {
@@ -171,6 +189,7 @@
                 }
                 mInfo.type = Display.TYPE_VIRTUAL;
                 mInfo.touch = DisplayDeviceInfo.TOUCH_NONE;
+                mInfo.state = mSurface != null ? Display.STATE_ON : Display.STATE_OFF;
                 mInfo.ownerUid = mOwnerUid;
                 mInfo.ownerPackageName = mOwnerPackageName;
             }