aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From kevi...@apache.org
Subject incubator-aurora git commit: Introduce AnnotatedAuroraAdmin superclass for annotations.
Date Tue, 17 Mar 2015 20:26:40 GMT
Repository: incubator-aurora
Updated Branches:
  refs/heads/master 8fc322940 -> 4b39d26f2


Introduce AnnotatedAuroraAdmin superclass for annotations.

Add an AnnotatedAuroraAdmin subclass for AuroraAdmin.Iface and change
code to inherit from it. This gives us a place to put annotations like
the AuthorizingParam one introduced in this review without needing to
copy-paste them when we override a new method. A future diff will use
these annotations to determine which permission a method call needs by
inspecting the annotated parameter. I created a new interface to enable
DRY - otherwise I'd need to annotate both ForwardingThrift and
SchedulerThriftInterface and keep them in sync.

Testing Done:
./gradlew build

Bugs closed: AURORA-1187

Reviewed at https://reviews.apache.org/r/32141/


Project: http://git-wip-us.apache.org/repos/asf/incubator-aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-aurora/commit/4b39d26f
Tree: http://git-wip-us.apache.org/repos/asf/incubator-aurora/tree/4b39d26f
Diff: http://git-wip-us.apache.org/repos/asf/incubator-aurora/diff/4b39d26f

Branch: refs/heads/master
Commit: 4b39d26f266c4274fbd2b50fb45ee09ff87b66ae
Parents: 8fc3229
Author: Kevin Sweeney <kevints@apache.org>
Authored: Tue Mar 17 13:26:33 2015 -0700
Committer: Kevin Sweeney <kevints@apache.org>
Committed: Tue Mar 17 13:26:33 2015 -0700

----------------------------------------------------------------------
 .../http/api/security/AuthorizingParam.java     |  72 ++++++++++
 .../thrift/SchedulerThriftInterface.java        |   4 +-
 .../aurora/scheduler/thrift/ThriftModule.java   |   2 +
 .../thrift/aop/AnnotatedAuroraAdmin.java        | 134 +++++++++++++++++++
 .../http/api/security/ApiSecurityIT.java        |   5 +-
 .../scheduler/thrift/aop/AopModuleTest.java     |   4 +-
 .../aop/FeatureToggleInterceptorTest.java       |   9 +-
 .../scheduler/thrift/aop/ForwardingThrift.java  |   7 +-
 .../thrift/aop/MockDecoratedThrift.java         |   8 +-
 .../thrift/aop/ServerInfoInterceptorTest.java   |   9 +-
 .../aop/ThriftStatsExporterInterceptorTest.java |   9 +-
 11 files changed, 235 insertions(+), 28 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/main/java/org/apache/aurora/scheduler/http/api/security/AuthorizingParam.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/http/api/security/AuthorizingParam.java
b/src/main/java/org/apache/aurora/scheduler/http/api/security/AuthorizingParam.java
new file mode 100644
index 0000000..8089879
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/http/api/security/AuthorizingParam.java
@@ -0,0 +1,72 @@
+/**
+ * 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.
+ */
+package org.apache.aurora.scheduler.http.api.security;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Signals to {@link org.apache.aurora.scheduler.http.api.security.ShiroThriftInterceptor}
that this
+ * parameter should be used to determine the instance the calling subject is attempting to
operate
+ * on. Methods using this parameter should ensure that the RPC cannot operate on an instance
+ * outside the scope of this parameter, otherwise a privilege escalation vulnerability exists.
+ *
+ * <p>
+ * A method intercepted by this interceptor that does not contain an AuthorizingParam or
with
+ * multiple AuthorizingParams will result in an {@link java.lang.IllegalStateException} from
the
+ * interceptor.
+ *
+ * <p>
+ * If the parameter type is not known to the interceptor an {@link java.lang.IllegalStateException}
+ * is thrown.
+ *
+ * <p>
+ * Inheritance semantics: while the Java Reflection API doesn't define inheritance for method
+ * parameter annotations the interceptor uses the first matching superclass where the method
is
+ * defined and annotated. This means that
+ *
+ * <pre>
+ * <code>
+ * public class A {
+ *   Response getItem(@AuthorizingParam String itemId) {
+ *     // ...
+ *   }
+ * }
+ * </code>
+ * </pre>
+ *
+ * and
+ *
+ * <pre>
+ * <code>
+ * public interface A {
+ *   Response getItem(@AuthorizingParam String itemId);
+ * }
+ *
+ * public class B implements A {
+ *   {@literal @}Override
+ *   Response getItem(String itemId) {
+ *     // ...
+ *   }
+ * }
+ * </code>
+ * </pre>
+ *
+ * are equivalent.
+ */
+@Target(ElementType.PARAMETER)
+@Retention(RetentionPolicy.RUNTIME)
+public @interface AuthorizingParam { }

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
index 6b15bfd..5588d17 100644
--- a/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/SchedulerThriftInterface.java
@@ -45,7 +45,6 @@ import org.apache.aurora.auth.CapabilityValidator.AuditCheck;
 import org.apache.aurora.auth.SessionValidator.AuthFailedException;
 import org.apache.aurora.gen.AcquireLockResult;
 import org.apache.aurora.gen.AddInstancesConfig;
-import org.apache.aurora.gen.AuroraAdmin;
 import org.apache.aurora.gen.ConfigRewrite;
 import org.apache.aurora.gen.DrainHostsResult;
 import org.apache.aurora.gen.EndMaintenanceResult;
@@ -121,6 +120,7 @@ import org.apache.aurora.scheduler.storage.entities.ILockKey;
 import org.apache.aurora.scheduler.storage.entities.IResourceAggregate;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
+import org.apache.aurora.scheduler.thrift.aop.AnnotatedAuroraAdmin;
 import org.apache.aurora.scheduler.thrift.auth.DecoratedThrift;
 import org.apache.aurora.scheduler.thrift.auth.Requires;
 import org.apache.aurora.scheduler.updater.JobDiff;
@@ -159,7 +159,7 @@ import static org.apache.aurora.scheduler.thrift.Responses.ok;
  * administration tasks.
  */
 @DecoratedThrift
-class SchedulerThriftInterface implements AuroraAdmin.Iface {
+class SchedulerThriftInterface implements AnnotatedAuroraAdmin {
   @Positive
   @CmdLine(name = "max_tasks_per_job", help = "Maximum number of allowed tasks in a single
job.")
   public static final Arg<Integer> MAX_TASKS_PER_JOB = Arg.create(4000);

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/main/java/org/apache/aurora/scheduler/thrift/ThriftModule.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/thrift/ThriftModule.java b/src/main/java/org/apache/aurora/scheduler/thrift/ThriftModule.java
index d9184eb..c65e021 100644
--- a/src/main/java/org/apache/aurora/scheduler/thrift/ThriftModule.java
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/ThriftModule.java
@@ -17,6 +17,7 @@ import com.google.inject.AbstractModule;
 
 import org.apache.aurora.gen.AuroraAdmin;
 import org.apache.aurora.gen.ReadOnlyScheduler;
+import org.apache.aurora.scheduler.thrift.aop.AnnotatedAuroraAdmin;
 import org.apache.aurora.scheduler.thrift.aop.AopModule;
 
 /**
@@ -28,6 +29,7 @@ public class ThriftModule extends AbstractModule {
   protected void configure() {
     bind(ReadOnlyScheduler.Iface.class).to(ReadOnlySchedulerImpl.class);
     bind(AuroraAdmin.Iface.class).to(SchedulerThriftInterface.class);
+    bind(AnnotatedAuroraAdmin.class).to(SchedulerThriftInterface.class);
 
     // Promote to an explicit binding so it's created in the servlet container child injector.
     // See https://code.google.com/p/google-guice/issues/detail?id=461

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/main/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdmin.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdmin.java
b/src/main/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdmin.java
new file mode 100644
index 0000000..6d30c25
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/thrift/aop/AnnotatedAuroraAdmin.java
@@ -0,0 +1,134 @@
+/**
+ * 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.
+ */
+package org.apache.aurora.scheduler.thrift.aop;
+
+import java.util.Set;
+
+import javax.annotation.Nullable;
+
+import org.apache.aurora.gen.AddInstancesConfig;
+import org.apache.aurora.gen.AuroraAdmin;
+import org.apache.aurora.gen.JobConfiguration;
+import org.apache.aurora.gen.JobKey;
+import org.apache.aurora.gen.JobUpdateKey;
+import org.apache.aurora.gen.JobUpdateRequest;
+import org.apache.aurora.gen.Lock;
+import org.apache.aurora.gen.LockKey;
+import org.apache.aurora.gen.LockValidation;
+import org.apache.aurora.gen.Response;
+import org.apache.aurora.gen.SessionKey;
+import org.apache.aurora.gen.TaskQuery;
+import org.apache.aurora.scheduler.http.api.security.AuthorizingParam;
+import org.apache.thrift.TException;
+
+/**
+ * Extension of the generated thrift interface with Java annotations, for example {@link
Nullable}
+ * and {@link AuthorizingParam}.
+ *
+ * When adding new methods to api.thrift you should add, at the very least, {@link Nullable}
+ * annotations for them here as well.
+ *
+ * TODO(ksweeney): Investigate adding other (JSR303) validation annotations here as well.
+ *
+ * TODO(ksweeney): Consider disallowing throwing subclasses by removing {@code throws TException}.
+ */
+public interface AnnotatedAuroraAdmin extends AuroraAdmin.Iface {
+  @Override
+  Response createJob(
+      @AuthorizingParam @Nullable JobConfiguration description,
+      @Nullable Lock lock,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response scheduleCronJob(
+      @AuthorizingParam @Nullable JobConfiguration description,
+      @Nullable Lock lock,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response descheduleCronJob(
+      @AuthorizingParam @Nullable JobKey job,
+      @Nullable Lock lock,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response startCronJob(
+      @AuthorizingParam @Nullable JobKey job,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response restartShards(
+      @AuthorizingParam @Nullable JobKey job,
+      @Nullable Set<Integer> shardIds,
+      @Nullable Lock lock,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response killTasks(
+      @AuthorizingParam @Nullable TaskQuery query,
+      @Nullable Lock lock,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response addInstances(
+      @AuthorizingParam @Nullable AddInstancesConfig config,
+      @Nullable Lock lock,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response acquireLock(
+      @AuthorizingParam @Nullable LockKey lockKey,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response releaseLock(
+      @AuthorizingParam @Nullable Lock lock,
+      @Nullable LockValidation validation,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response replaceCronTemplate(
+      @AuthorizingParam @Nullable JobConfiguration config,
+      @Nullable Lock lock,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response startJobUpdate(
+      @AuthorizingParam @Nullable JobUpdateRequest request,
+      @Nullable String message,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response pauseJobUpdate(
+      @AuthorizingParam @Nullable JobUpdateKey key,
+      @Nullable String message,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response resumeJobUpdate(
+      @AuthorizingParam @Nullable JobUpdateKey key,
+      @Nullable String message,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response abortJobUpdate(
+      @AuthorizingParam @Nullable JobUpdateKey key,
+      @Nullable String message,
+      @Nullable SessionKey session) throws TException;
+
+  @Override
+  Response pulseJobUpdate(
+      @AuthorizingParam @Nullable JobUpdateKey key,
+      @Nullable SessionKey session) throws TException;
+}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/test/java/org/apache/aurora/scheduler/http/api/security/ApiSecurityIT.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/http/api/security/ApiSecurityIT.java
b/src/test/java/org/apache/aurora/scheduler/http/api/security/ApiSecurityIT.java
index f0e40b6..76cb691 100644
--- a/src/test/java/org/apache/aurora/scheduler/http/api/security/ApiSecurityIT.java
+++ b/src/test/java/org/apache/aurora/scheduler/http/api/security/ApiSecurityIT.java
@@ -31,6 +31,7 @@ import org.apache.aurora.gen.Response;
 import org.apache.aurora.gen.ResponseCode;
 import org.apache.aurora.scheduler.http.JettyServerModuleTest;
 import org.apache.aurora.scheduler.http.api.ApiModule;
+import org.apache.aurora.scheduler.thrift.aop.AnnotatedAuroraAdmin;
 import org.apache.aurora.scheduler.thrift.aop.MockDecoratedThrift;
 import org.apache.http.auth.AuthScope;
 import org.apache.http.auth.Credentials;
@@ -79,7 +80,7 @@ public class ApiSecurityIT extends JettyServerModuleTest {
       ImmutableSet.<Credentials>of(ROOT, WFARNER, UNPRIVILEGED, BACKUP_SERVICE);
 
   private Ini ini;
-  private AuroraAdmin.Iface auroraAdmin;
+  private AnnotatedAuroraAdmin auroraAdmin;
   private StatsProvider statsProvider;
 
   @Before
@@ -97,7 +98,7 @@ public class ApiSecurityIT extends JettyServerModuleTest {
     roles.put("eng", "thrift.AuroraSchedulerManager:*");
     roles.put("backupsvc", "thrift.AuroraAdmin:listBackups");
 
-    auroraAdmin = createMock(AuroraAdmin.Iface.class);
+    auroraAdmin = createMock(AnnotatedAuroraAdmin.class);
     statsProvider = createMock(StatsProvider.class);
     expect(statsProvider.makeCounter(anyString())).andStubReturn(new AtomicLong());
   }

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/test/java/org/apache/aurora/scheduler/thrift/aop/AopModuleTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/aop/AopModuleTest.java b/src/test/java/org/apache/aurora/scheduler/thrift/aop/AopModuleTest.java
index df6b53a..d20c9da 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/aop/AopModuleTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/aop/AopModuleTest.java
@@ -43,7 +43,7 @@ public class AopModuleTest extends EasyMockTest {
   private static final SessionKey SESSION_KEY = new SessionKey();
 
   private CapabilityValidator capabilityValidator;
-  private Iface mockThrift;
+  private AnnotatedAuroraAdmin mockThrift;
 
   @Before
   public void setUp() throws Exception {
@@ -51,7 +51,7 @@ public class AopModuleTest extends EasyMockTest {
     expect(capabilityValidator.toString(SESSION_KEY))
         .andReturn("user")
         .anyTimes();
-    mockThrift = createMock(Iface.class);
+    mockThrift = createMock(AnnotatedAuroraAdmin.class);
   }
 
   private Iface getIface(Map<String, Boolean> toggledMethods) {

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/test/java/org/apache/aurora/scheduler/thrift/aop/FeatureToggleInterceptorTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/aop/FeatureToggleInterceptorTest.java
b/src/test/java/org/apache/aurora/scheduler/thrift/aop/FeatureToggleInterceptorTest.java
index ee98f66..22a1546 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/aop/FeatureToggleInterceptorTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/aop/FeatureToggleInterceptorTest.java
@@ -23,7 +23,6 @@ import com.google.inject.TypeLiteral;
 import com.google.inject.matcher.Matchers;
 import com.twitter.common.testing.easymock.EasyMockTest;
 
-import org.apache.aurora.gen.AuroraAdmin;
 import org.apache.aurora.gen.Response;
 import org.apache.aurora.gen.ResponseCode;
 import org.apache.aurora.gen.TaskQuery;
@@ -37,13 +36,13 @@ import static org.junit.Assert.assertSame;
 
 public class FeatureToggleInterceptorTest extends EasyMockTest {
 
-  private AuroraAdmin.Iface realThrift;
-  private AuroraAdmin.Iface decoratedThrift;
+  private AnnotatedAuroraAdmin realThrift;
+  private AnnotatedAuroraAdmin decoratedThrift;
   private Predicate<Method> predicate;
 
   @Before
   public void setUp() {
-    realThrift = createMock(AuroraAdmin.Iface.class);
+    realThrift = createMock(AnnotatedAuroraAdmin.class);
     predicate = createMock(new Clazz<Predicate<Method>>() { });
     Injector injector = Guice.createInjector(new AbstractModule() {
       @Override
@@ -56,7 +55,7 @@ public class FeatureToggleInterceptorTest extends EasyMockTest {
             new FeatureToggleInterceptor());
       }
     });
-    decoratedThrift = injector.getInstance(AuroraAdmin.Iface.class);
+    decoratedThrift = injector.getInstance(AnnotatedAuroraAdmin.class);
   }
 
   @Test

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/test/java/org/apache/aurora/scheduler/thrift/aop/ForwardingThrift.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/aop/ForwardingThrift.java b/src/test/java/org/apache/aurora/scheduler/thrift/aop/ForwardingThrift.java
index 3900c22..d25618c 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/aop/ForwardingThrift.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/aop/ForwardingThrift.java
@@ -16,7 +16,6 @@ package org.apache.aurora.scheduler.thrift.aop;
 import java.util.Set;
 
 import org.apache.aurora.gen.AddInstancesConfig;
-import org.apache.aurora.gen.AuroraAdmin;
 import org.apache.aurora.gen.Hosts;
 import org.apache.aurora.gen.JobConfiguration;
 import org.apache.aurora.gen.JobKey;
@@ -40,11 +39,11 @@ import static java.util.Objects.requireNonNull;
  * A forwarding scheduler controller to make it easy to override specific behavior in an
  * implementation class.
  */
-abstract class ForwardingThrift implements AuroraAdmin.Iface {
+abstract class ForwardingThrift implements AnnotatedAuroraAdmin {
 
-  private final AuroraAdmin.Iface delegate;
+  private final AnnotatedAuroraAdmin delegate;
 
-  ForwardingThrift(AuroraAdmin.Iface delegate) {
+  ForwardingThrift(AnnotatedAuroraAdmin delegate) {
     this.delegate = requireNonNull(delegate);
   }
 

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/test/java/org/apache/aurora/scheduler/thrift/aop/MockDecoratedThrift.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/aop/MockDecoratedThrift.java
b/src/test/java/org/apache/aurora/scheduler/thrift/aop/MockDecoratedThrift.java
index 7b1bf2e..1415f0c 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/aop/MockDecoratedThrift.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/aop/MockDecoratedThrift.java
@@ -41,12 +41,14 @@ public class MockDecoratedThrift extends ForwardingThrift {
   private @interface MockThrift { }
 
   @Inject
-  MockDecoratedThrift(@MockThrift AuroraAdmin.Iface delegate) {
+  MockDecoratedThrift(@MockThrift AnnotatedAuroraAdmin delegate) {
     super(delegate);
   }
 
-  public static void bindForwardedMock(Binder binder, AuroraAdmin.Iface mockThrift) {
-    binder.bind(AuroraAdmin.Iface.class).annotatedWith(MockThrift.class).toInstance(mockThrift);
+  public static void bindForwardedMock(Binder binder, AnnotatedAuroraAdmin mockThrift) {
+    binder.bind(AnnotatedAuroraAdmin.class).annotatedWith(MockThrift.class).toInstance(mockThrift);
+
+    binder.bind(AnnotatedAuroraAdmin.class).to(MockDecoratedThrift.class);
     binder.bind(AuroraAdmin.Iface.class).to(MockDecoratedThrift.class);
   }
 }

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/test/java/org/apache/aurora/scheduler/thrift/aop/ServerInfoInterceptorTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/aop/ServerInfoInterceptorTest.java
b/src/test/java/org/apache/aurora/scheduler/thrift/aop/ServerInfoInterceptorTest.java
index 7e20aaa..793684c 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/aop/ServerInfoInterceptorTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/aop/ServerInfoInterceptorTest.java
@@ -20,7 +20,6 @@ import com.google.inject.Injector;
 import com.google.inject.matcher.Matchers;
 import com.twitter.common.testing.easymock.EasyMockTest;
 
-import org.apache.aurora.gen.AuroraAdmin;
 import org.apache.aurora.gen.GetJobsResult;
 import org.apache.aurora.gen.JobConfiguration;
 import org.apache.aurora.gen.Response;
@@ -39,8 +38,8 @@ public class ServerInfoInterceptorTest extends EasyMockTest {
 
   private static final String ROLE = "bob";
 
-  private AuroraAdmin.Iface realThrift;
-  private AuroraAdmin.Iface decoratedThrift;
+  private AnnotatedAuroraAdmin realThrift;
+  private AnnotatedAuroraAdmin decoratedThrift;
 
   private static final IServerInfo SERVER_INFO = IServerInfo.build(
       new ServerInfo()
@@ -53,7 +52,7 @@ public class ServerInfoInterceptorTest extends EasyMockTest {
   @Before
   public void setUp() {
     interceptor = new ServerInfoInterceptor();
-    realThrift = createMock(AuroraAdmin.Iface.class);
+    realThrift = createMock(AnnotatedAuroraAdmin.class);
     Injector injector = Guice.createInjector(new AbstractModule() {
       @Override
       protected void configure() {
@@ -65,7 +64,7 @@ public class ServerInfoInterceptorTest extends EasyMockTest {
             interceptor);
       }
     });
-    decoratedThrift = injector.getInstance(AuroraAdmin.Iface.class);
+    decoratedThrift = injector.getInstance(AnnotatedAuroraAdmin.class);
   }
 
   @Test

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/4b39d26f/src/test/java/org/apache/aurora/scheduler/thrift/aop/ThriftStatsExporterInterceptorTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/thrift/aop/ThriftStatsExporterInterceptorTest.java
b/src/test/java/org/apache/aurora/scheduler/thrift/aop/ThriftStatsExporterInterceptorTest.java
index aae5cd7..dac9af2 100644
--- a/src/test/java/org/apache/aurora/scheduler/thrift/aop/ThriftStatsExporterInterceptorTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/thrift/aop/ThriftStatsExporterInterceptorTest.java
@@ -21,7 +21,6 @@ import com.google.inject.matcher.Matchers;
 import com.twitter.common.stats.Stats;
 import com.twitter.common.testing.easymock.EasyMockTest;
 
-import org.apache.aurora.gen.AuroraAdmin;
 import org.apache.aurora.gen.GetJobsResult;
 import org.apache.aurora.gen.JobConfiguration;
 import org.apache.aurora.gen.Response;
@@ -39,14 +38,14 @@ public class ThriftStatsExporterInterceptorTest extends EasyMockTest {
 
   private static final String ROLE = "bob";
 
-  private AuroraAdmin.Iface realThrift;
-  private AuroraAdmin.Iface decoratedThrift;
+  private AnnotatedAuroraAdmin realThrift;
+  private AnnotatedAuroraAdmin decoratedThrift;
   private ThriftStatsExporterInterceptor statsInterceptor;
 
   @Before
   public void setUp() {
     statsInterceptor = new ThriftStatsExporterInterceptor();
-    realThrift = createMock(AuroraAdmin.Iface.class);
+    realThrift = createMock(AnnotatedAuroraAdmin.class);
     Injector injector = Guice.createInjector(new AbstractModule() {
       @Override
       protected void configure() {
@@ -57,7 +56,7 @@ public class ThriftStatsExporterInterceptorTest extends EasyMockTest {
             statsInterceptor);
       }
     });
-    decoratedThrift = injector.getInstance(AuroraAdmin.Iface.class);
+    decoratedThrift = injector.getInstance(AnnotatedAuroraAdmin.class);
   }
 
   @Test


Mime
View raw message