Add cancellation of protection process

yck1509 · Jun 22, 2014 · 5969669 · 5969669
1 parent 010ae7d
commit 5969669
Show file tree

Hide file tree

Showing 15 changed files with 74 additions and 7 deletions.
diff --git a/Confuser.Core/ConfuserEngine.cs b/Confuser.Core/ConfuserEngine.cs
@@ -66,7 +66,6 @@ private static void RunInternal(ConfuserParameters parameters, CancellationToken
 			context.PackerInitiated = parameters.PackerInitiated;
 			context.token = token;
 
-
 			PrintInfo(context);
 
 			bool ok = false;
@@ -80,6 +79,8 @@ private static void RunInternal(ConfuserParameters parameters, CancellationToken
 				foreach (string probePath in parameters.Project.ProbePaths)
 					asmResolver.PostSearchPaths.Add(Path.Combine(context.BaseDirectory, probePath));
 
+				context.CheckCancellation();
+
 				Marker marker = parameters.GetMarker();
 
 				// 2. Discover plugins
@@ -92,6 +93,8 @@ private static void RunInternal(ConfuserParameters parameters, CancellationToken
 
 				context.Logger.InfoFormat("Discovered {0} protections, {1} packers.", prots.Count, packers.Count);
 
+				context.CheckCancellation();
+
 				// 3. Resolve dependency
 				context.Logger.Debug("Resolving component dependency...");
 				try {
@@ -108,6 +111,8 @@ private static void RunInternal(ConfuserParameters parameters, CancellationToken
 				foreach (Packer packer in packers)
 					components.Add(packer);
 
+				context.CheckCancellation();
+
 				// 4. Load modules
 				context.Logger.Info("Loading input modules...");
 				marker.Initalize(prots, packers);
@@ -119,6 +124,8 @@ private static void RunInternal(ConfuserParameters parameters, CancellationToken
 					asmResolver.AddToCache(module);
 				context.Packer = markings.Packer;
 
+				context.CheckCancellation();
+
 				// 5. Initialize components
 				context.Logger.Info("Initializing...");
 				foreach (ConfuserComponent comp in components) {
@@ -131,13 +138,16 @@ private static void RunInternal(ConfuserParameters parameters, CancellationToken
 					context.CheckCancellation();
 				}
 
+				context.CheckCancellation();
+
 				// 6. Build pipeline
 				context.Logger.Debug("Building pipeline...");
 				var pipeline = new ProtectionPipeline();
 				context.Pipeline = pipeline;
 				foreach (ConfuserComponent comp in components) {
 					comp.PopulatePipeline(pipeline);
 				}
+
 				context.CheckCancellation();
 
 				//7. Run pipeline
@@ -285,6 +295,7 @@ private static void BeginModule(ConfuserContext context) {
 			context.Logger.InfoFormat("Processing module '{0}'...", context.CurrentModule.Name);
 
 			context.CurrentModuleWriterListener = new ModuleWriterListener();
+			context.CurrentModuleWriterListener.OnWriterEvent += (sender, e) => context.CheckCancellation();
 			context.CurrentModuleWriterOptions = new ModuleWriterOptions(context.CurrentModule, context.CurrentModuleWriterListener);
 			var snKey = context.Annotations.Get<StrongNameKey>(context.CurrentModule, Marker.SNKey);
 			context.CurrentModuleWriterOptions.InitializeStrongNameSigning(context.CurrentModule, snKey);

diff --git a/Confuser.Core/Marker.cs b/Confuser.Core/Marker.cs
@@ -105,13 +105,17 @@ protected internal virtual MarkerResult MarkProject(ConfuserProject proj, Confus
 			foreach (ProjectModule module in proj) {
 				context.Logger.InfoFormat("Loading '{0}'...", module.Path);
 				ModuleDefMD modDef = module.Resolve(proj.BaseDirectory, context.Resolver.DefaultModuleContext);
+				context.CheckCancellation();
+
 				Rules rules = ParseRules(proj, module, context);
 
 				context.Annotations.Set(modDef, SNKey, LoadSNKey(context, module.SNKeyPath == null ? null : Path.Combine(proj.BaseDirectory, module.SNKeyPath), module.SNKeyPassword));
 				context.Annotations.Set(modDef, RulesKey, rules);
 
-				foreach (IDnlibDef def in modDef.FindDefinitions())
+				foreach (IDnlibDef def in modDef.FindDefinitions()) {
 					ApplyRules(context, def, rules);
+					context.CheckCancellation();
+				}
 
 				// Packer parameters are stored in modules
 				if (packerParams != null)

diff --git a/Confuser.Core/ProtectionPipeline.cs b/Confuser.Core/ProtectionPipeline.cs
@@ -120,10 +120,11 @@ public T FindPhase<T>() where T : ProtectionPhase {
 		/// <param name="context">The working context.</param>
 		internal void ExecuteStage(PipelineStage stage, Action<ConfuserContext> func, Func<IList<IDnlibDef>> targets, ConfuserContext context) {
 			foreach (ProtectionPhase pre in preStage[stage]) {
+				context.CheckCancellation();
 				context.Logger.DebugFormat("Executing '{0}' phase...", pre.Name);
 				pre.Execute(context, new ProtectionParameters(pre.Parent, Filter(context, targets(), pre)));
-				context.CheckCancellation();
 			}
+			context.CheckCancellation();
 			func(context);
 			context.CheckCancellation();
 			foreach (ProtectionPhase post in postStage[stage]) {

diff --git a/Confuser.Protections/AntiTamper/JITMode.cs b/Confuser.Protections/AntiTamper/JITMode.cs
@@ -218,6 +218,7 @@ private void CreateSection(ModuleWriter writer) {
 
 				method.Body = NopBody;
 				writer.MetaData.TablesHeap.MethodTable[token.Rid].ImplFlags |= (ushort)MethodImplAttributes.NoInlining;
+				context.CheckCancellation();
 			}
 			bodyIndex.PopulateSection(newSection);
 

diff --git a/Confuser.Protections/Compress/Compressor.cs b/Confuser.Protections/Compress/Compressor.cs
@@ -77,6 +77,7 @@ protected override void Pack(ConfuserContext context, ProtectionParameters param
 				stubModule.Write(ms, new ModuleWriterOptions(stubModule, new KeyInjector(ctx)) {
 					StrongNameKey = snKey
 				});
+				context.CheckCancellation();
 				base.ProtectStub(context, context.OutputPaths[ctx.ModuleIndex], ms.ToArray(), snKey, new StubProtection(ctx));
 			}
 		}
@@ -118,6 +119,7 @@ private void PackModules(ConfuserContext context, CompressorContext compCtx, Mod
 					progress = (progress + moduleIndex) / modules.Count;
 					context.Logger.Progress((int)(progress * 10000), 10000);
 				});
+				context.CheckCancellation();
 
 				var resource = new EmbeddedResource(Convert.ToBase64String(name), encrypted, ManifestResourceAttributes.Private);
 				stubModule.Resources.Add(resource);
@@ -183,6 +185,7 @@ private void InjectStub(ConfuserContext context, CompressorContext compCtx, Prot
 			byte[] encryptedModule = compCtx.Encrypt(comp, compCtx.OriginModule, seed,
 			                                         progress => context.Logger.Progress((int)(progress * 10000), 10000));
 			context.Logger.EndProgress();
+			context.CheckCancellation();
 
 			compCtx.EncryptedModule = encryptedModule;
 

diff --git a/Confuser.Protections/Constants/EncodePhase.cs b/Confuser.Protections/Constants/EncodePhase.cs
@@ -40,6 +40,7 @@ protected override void Execute(ConfuserContext context, ProtectionParameters pa
 			foreach (var entry in ldInit.WithProgress(context.Logger)) // Ensure the array length haven't been encoded yet
 			{
 				EncodeInitializer(moduleCtx, entry.Key, entry.Value);
+				context.CheckCancellation();
 			}
 			foreach (var entry in ldc.WithProgress(context.Logger)) {
 				if (entry.Key is string) {
@@ -63,6 +64,7 @@ protected override void Execute(ConfuserContext context, ProtectionParameters pa
 				}
 				else
 					throw new UnreachableException();
+				context.CheckCancellation();
 			}
 
 			// compress
@@ -76,6 +78,7 @@ protected override void Execute(ConfuserContext context, ProtectionParameters pa
 			}
 			Debug.Assert(buffIndex == encodedBuff.Length);
 			encodedBuff = context.Registry.GetService<ICompressionService>().Compress(encodedBuff);
+			context.CheckCancellation();
 
 			uint compressedLen = (uint)(encodedBuff.Length + 3) / 4;
 			compressedLen = (compressedLen + 0xfu) & ~0xfu;
@@ -323,6 +326,8 @@ private void ExtractConstants(
 					if (eligible)
 						ldc.AddListEntry(instr.Operand, Tuple.Create(method, instr));
 				}
+
+				context.CheckCancellation();
 			}
 		}
 

diff --git a/Confuser.Protections/ControlFlow/ControlFlowPhase.cs b/Confuser.Protections/ControlFlow/ControlFlowPhase.cs
@@ -74,6 +74,7 @@ protected override void Execute(ConfuserContext context, ProtectionParameters pa
 			foreach (MethodDef method in parameters.Targets.OfType<MethodDef>().WithProgress(context.Logger))
 				if (method.HasBody && method.Body.Instructions.Count > 0) {
 					ProcessMethod(method.Body, ParseParameters(method, context, parameters, random, disabledOpti));
+					context.CheckCancellation();
 				}
 		}
 

diff --git a/Confuser.Protections/ReferenceProxy/ReferenceProxyPhase.cs b/Confuser.Protections/ReferenceProxy/ReferenceProxyPhase.cs
@@ -100,6 +100,7 @@ protected override void Execute(ConfuserContext context, ProtectionParameters pa
 			foreach (MethodDef method in parameters.Targets.OfType<MethodDef>().WithProgress(context.Logger))
 				if (method.HasBody && method.Body.Instructions.Count > 0) {
 					ProcessMethod(ParseParameters(method, context, parameters, store));
+					context.CheckCancellation();
 				}
 
 			RPContext ctx = ParseParameters(context.CurrentModule, context, parameters, store);

diff --git a/Confuser.Protections/Resources/MDPhase.cs b/Confuser.Protections/Resources/MDPhase.cs
@@ -26,6 +26,7 @@ public void Hook() {
 		private void OnWriterEvent(object sender, ModuleWriterListenerEventArgs e) {
 			var writer = (ModuleWriter)sender;
 			if (e.WriterEvent == ModuleWriterEvent.MDBeginAddResources) {
+				ctx.Context.CheckCancellation();
 				ctx.Context.Logger.Debug("Encrypting resources...");
 
 				List<EmbeddedResource> resources = ctx.Module.Resources.OfType<EmbeddedResource>().ToList();
@@ -57,6 +58,7 @@ private void OnWriterEvent(object sender, ModuleWriterListenerEventArgs e) {
 					moduleBuff,
 					progress => ctx.Context.Logger.Progress((int)(progress * 10000), 10000));
 				ctx.Context.Logger.EndProgress();
+				ctx.Context.CheckCancellation();
 
 				uint compressedLen = (uint)(moduleBuff.Length + 3) / 4;
 				compressedLen = (compressedLen + 0xfu) & ~0xfu;

diff --git a/Confuser.Renamer/AnalyzePhase.cs b/Confuser.Renamer/AnalyzePhase.cs
@@ -43,12 +43,14 @@ protected override void Execute(ConfuserContext context, ProtectionParameters pa
 					service.GetVTables().GetVTable((TypeDef)def);
 					service.SetOriginalNamespace(def, ((TypeDef)def).Namespace);
 				}
+				context.CheckCancellation();
 			}
 
 			context.Logger.Debug("Analyzing...");
 			IList<IRenamer> renamers = service.Renamers;
 			foreach (IDnlibDef def in parameters.Targets.WithProgress(context.Logger)) {
 				Analyze(service, context, def, true);
+				context.CheckCancellation();
 			}
 		}
 

diff --git a/Confuser.Renamer/PostRenamePhase.cs b/Confuser.Renamer/PostRenamePhase.cs
@@ -24,6 +24,7 @@ protected override void Execute(ConfuserContext context, ProtectionParameters pa
 			foreach (IRenamer renamer in service.Renamers) {
 				foreach (IDnlibDef def in parameters.Targets)
 					renamer.PostRename(context, service, def);
+				context.CheckCancellation();
 			}
 		}
 	}

diff --git a/Confuser.Renamer/RenamePhase.cs b/Confuser.Renamer/RenamePhase.cs
@@ -23,6 +23,7 @@ protected override void Execute(ConfuserContext context, ProtectionParameters pa
 			foreach (IRenamer renamer in service.Renamers) {
 				foreach (IDnlibDef def in parameters.Targets)
 					renamer.PreRename(context, service, def);
+				context.CheckCancellation();
 			}
 
 			foreach (IDnlibDef def in parameters.Targets.WithProgress(context.Logger)) {
@@ -66,6 +67,7 @@ protected override void Execute(ConfuserContext context, ProtectionParameters pa
 						throw new ConfuserException(null);
 					}
 				}
+				context.CheckCancellation();
 			}
 		}
 	}

diff --git a/ConfuserEx/ConfuserEx.csproj b/ConfuserEx/ConfuserEx.csproj
@@ -228,4 +228,20 @@
   <Target Name="AfterBuild">
   </Target>
   -->
+
+  <Target Name="solveAliasProblem" >
+    <ConvertToAbsolutePath Paths="..\packages\TaskParallelLibrary.1.0.2856.0\lib\Net35\System.Threading.dll">
+      <Output TaskParameter="AbsolutePaths" PropertyName="FullPathOriginal" />
+    </ConvertToAbsolutePath>
+    <ItemGroup>
+      <ReferencePath Remove="$(FullPathOriginal)"/>
+      <ReferencePath Include="$(FullPathOriginal)">
+        <Aliases>PTL</Aliases>
+      </ReferencePath>
+    </ItemGroup>
+  </Target>
+  <PropertyGroup>
+    <CoreCompileDependsOn>solveAliasProblem;$(PrepareResourcesDependsOn)</CoreCompileDependsOn>
+  </PropertyGroup>
+
 </Project>
diff --git a/ConfuserEx/ViewModel/UI/ProtectTabVM.cs b/ConfuserEx/ViewModel/UI/ProtectTabVM.cs
@@ -1,11 +1,16 @@
-using System;
+extern alias PTL;
+
+using System;
 using System.Windows;
 using System.Windows.Documents;
 using System.Windows.Input;
 using System.Windows.Media;
 using Confuser.Core;
 using Confuser.Core.Project;
 using GalaSoft.MvvmLight.Command;
+using PTL::System.Threading;
+
+// http://connect.microsoft.com/VisualStudio/feedback/details/615953/
 
 namespace ConfuserEx.ViewModel {
 	internal class ProtectTabVM : TabViewModel, ILogger {
@@ -24,6 +29,10 @@ public ICommand ProtectCmd {
 			get { return new RelayCommand(DoProtect, () => !App.NavigationDisabled); }
 		}
 
+		public ICommand CancelCmd {
+			get { return new RelayCommand(DoCancel, () => App.NavigationDisabled); }
+		}
+
 		public double? Progress {
 			get { return progress; }
 			set { SetProperty(ref progress, value, "Progress"); }
@@ -36,18 +45,20 @@ public bool? Result {
 			set { SetProperty(ref result, value, "Result"); }
 		}
 
+		CancellationTokenSource cancelSrc;
 		private void DoProtect() {
 			var parameters = new ConfuserParameters();
 			parameters.Project = ((IViewModel<ConfuserProject>)App.Project).Model;
 			parameters.Logger = this;
 
 			documentContent.Inlines.Clear();
-			App.NavigationDisabled = true;
+			cancelSrc = new CancellationTokenSource();
 			Result = null;
 			Progress = null;
 			begin = DateTime.Now;
+			App.NavigationDisabled = true;
 
-			ConfuserEngine.Run(parameters)
+			ConfuserEngine.Run(parameters, cancelSrc.Token)
 			              .ContinueWith(_ =>
 			                            Application.Current.Dispatcher.BeginInvoke(new Action(() => {
 				                            Progress = 0;
@@ -56,6 +67,10 @@ private void DoProtect() {
 			                            })));
 		}
 
+		private void DoCancel() {
+			cancelSrc.Cancel();
+		}
+
 		private void AppendLine(string format, Brush foreground, params object[] args) {
 			Application.Current.Dispatcher.BeginInvoke(new Action(() => {
 				documentContent.Inlines.Add(new Run(string.Format(format, args)) { Foreground = foreground });

diff --git a/ConfuserEx/Views/ProtectTabView.xaml b/ConfuserEx/Views/ProtectTabView.xaml
@@ -11,12 +11,14 @@
             <Grid.ColumnDefinitions>
                 <ColumnDefinition Width="*" />
                 <ColumnDefinition Width="100px" />
+                <ColumnDefinition Width="100px" />
             </Grid.ColumnDefinitions>
 
             <ProgressBar x:Name="progress" Grid.Row="0" Grid.Column="0" Margin="5"
                          Value="{Binding Progress}" Minimum="0" Maximum="1" />
             <Button Grid.Row="0" Grid.Column="1" Margin="5" Content="Protect!" Command="{Binding ProtectCmd}" />
-            <RichTextBox x:Name="log" Grid.Row="1" Grid.ColumnSpan="2" Margin="5" FontFamily="Consolas"
+            <Button Grid.Row="0" Grid.Column="2" Margin="5" Content="Cancel" Command="{Binding CancelCmd}" />
+            <RichTextBox x:Name="log" Grid.Row="1" Grid.ColumnSpan="3" Margin="5" FontFamily="Consolas"
                          IsReadOnly="True" IsReadOnlyCaretVisible="True" local:Skin.RTBDocument="{Binding LogDocument}"
                          VerticalScrollBarVisibility="Visible" />
         </Grid>